Vector valued gradients?

[mathDR]

I think I must be totally missing something obvious when constructing custom jacobian vector products. A minimal example follows:

Assume we want to write a custom derivative for a quadratic form:

from functools import partial
import jax
import jax.numpy as jnp

@partial(jax.custom_jvp, nondiff_argnums=(1,))
def f(x, A):
    return jnp.dot(x,jnp.dot(A,x))

@f.defjvp
def f_jvp(A,primals,tangents):
    x, = primals
    x_dot, = tangents
    return jnp.dot(x,jnp.dot(A,x)), jnp.dot(A+A.T,x) * x_dot

and using some data:

x = jnp.array([0.5,0.25])
A = jnp.array([[1.25,1.0],[-0.5,0.25]])

computing the function is fine:

f(x,A) = 0.390625

but attempting to compute the gradient

print(jax.grad(f)(x,A))

yields the following error:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[83], line 1
----> 1 print(jax.grad(f)(x,A))

    [... skipping hidden 15 frame]

File ~/.pyenv/versions/3.11.1/envs/stan/lib/python3.11/site-packages/jax/_src/custom_derivatives.py:342, in _flatten_jvp(primal_name, jvp_name, in_tree, maybe_out_type, *args)
    339   (av1,), (av2,) = primal_avals_out, tangent_avals_out
    340   msg = ("Custom JVP rule must produce primal and tangent outputs with "
    341          "equal shapes and dtypes, but got {} and {} respectively.")
--> 342   raise TypeError(msg.format(av1.str_short(), av2.str_short()))
    343 else:
    344   msg = ("Custom JVP rule must produce primal and tangent outputs with "
    345          "equal shapes and dtypes, but got:\n{}")

TypeError: Custom JVP rule must produce primal and tangent outputs with equal shapes and dtypes, but got float32[] and float32[2] respectively.

How can this be reconciled? The original function is from $\mathbb{R}^2\rightarrow\mathbb{R}$, so it's gradient must be from $\mathbb{R}^2\rightarrow\mathbb{R}^2$. What am I missing/doing wrong?

Thanks!

[mathDR]

I will go ahead and answer this myself. It is because it is called a jacobian vector product. 🤦‍♂️

So the correct function is given by:

from functools import partial
import jax
import jax.numpy as jnp

@partial(jax.custom_jvp, nondiff_argnums=(1,))
def f(x, A):
    return jnp.dot(x,jnp.dot(A,x))

@f.defjvp
def f_jvp(A,primals,tangents):
    x, = primals
    x_dot, = tangents
    return jnp.dot(x,jnp.dot(A,x)), jnp.dot(jnp.dot(A+A.T,x),x_dot)