What is the "JAX" way of defining elementwise vector functions

[ipcamit]

Hi, Sorry if this question is already answered (I am sure it must be), but somehow I could not find any resource.
I have a $\R_2 \rightarrow \R_3 $ function like:

def fv(x):
     return np.array([
     x[0] + x[1],
     x[0] * x[1],
     np.log(x[0])*x[1]])

I can differentiate it explicitly using vjp. But I was wondering is there is any way to make vmap(grad(fv))(x) work? What is the recommended way to define such function transforms?

I get the following error:

IndexError                                Traceback (most recent call last)
Cell In [85], line 1
----> 1 jax.vmap(jax.grad(fv))(x)
Cell In [72], line 3, in fv(x)
      1 def fv(x):
      2     return np.array([
----> 3     x[0] + x[1],
...
-> 4402   raise IndexError(
   4403       f"Too many indices for {array_name}: {len_without_none} "
...
IndexError: Too many indices for array: 1 non-None/Ellipsis indices for dim 0.

[jakevdp]

There's no way to make vmap(grad(fv)) work for this function, because grad requires a scalar-output function, and your fv function is a vector output that does not correspond to any vector input. If you're interested on the gradient's effect on a single output, you could do something like this:

jax.grad(lambda x: fv(x)[0])(x)

If you're interested in computing the gradients for all outputs at once, you could do something like this:

jax.jacrev(fv)(x)