vmap + jacfwd error when attempting to compute higher order derivatives of a 1D at an array of values

[jdongg]

A simple MWE is included below, which computes up to the third derivative of sin(x) for an array of values of x. I accomplish this with jacfwd and vmap. There are no issues up to the second derivative and the results when plotted appear correct. I get this error with the third derivative:

df3 = jnp.reshape(dfdx3(x), x.shape)
ValueError: vmap was requested to map its argument along axis 0, which implies that its rank should be at least 1, but is only 0 (its shape is ())

and am a bit unsure what to make of it. It says dfdx2 has no shape when I would expect it to be N x 1 x 1 (which is confirmed when dfdx2 is evaluated at x... am I going about computing the derivatives in a way that doesn't make sense in jax?

import jax.numpy as jnp
import numpy as np

def mySin(x):
	return jnp.sin(x)

x = np.reshape(np.linspace(0., 2.*np.pi, 100), [100, 1])
f = mySin(x)
dfdx = jax.vmap(jax.jacfwd(mySin, 0), 0, 0)
dfdx2 = jax.vmap(jax.jacfwd(dfdx, 0), 0, 0)
dfdx3 = jax.vmap(jax.jacfwd(dfdx2, 0), 0, 0)

df = jnp.reshape(dfdx(x), x.shape)
df2 = jnp.reshape(dfdx2(x), x.shape)
df3 = jnp.reshape(dfdx3(x), x.shape)

[jakevdp]

vmap refers to the axis of the input arrays, not to the function that is mapped. So the relevant axes here are not anything associated with dfdx2, but rather the axes of the input argument x.

The input argument is x, which is of shape [100, 1]. The first vmap maps over the first axis of length 100; the second nested vmap maps over the second axis of length 1; the third nested vmap doesn't have any axis to map over, therefore it raises this error.

From your example, it's not entirely clear to me what you're trying to compute (what exactly is a triply nested vmapped jacobian? I'm honestly not sure). That said, if you want to compute the element-wise third derivative of a scalar function applied to a 2D input, one option might look like this:

from jax import vmap, grad
df3 = grad(grad(grad(mySin)))
out = vmap(vmap(df3))(x)
print(out.shape)
# (100, 1)

We use grad three times because grad is appropriate for scalar outputs, and the function maps a scalar to a scalar. We then use vmap twice to map over for the two axes of the input x.

Does that answer your question?