[Python] Is there a convenient way to take the gradient/jacobian of an already vmapped...

There is a function in my code base that is "already vmapped", i.e. when fed an array of shape (M, N) it outputs another array of shape (M, N). I would like to take the "row-wise Jacobian" of this function: i.e. a function that returns an array of shape (M, N, N). I've achieved this so far in hacky way by adding a dummy extra dimension and vmapping, as illustrated in the example below, but it really feels like there should be a better way to do this. Does anyone have any ideas?

Example of what I want:

import jax
from jax import numpy as jnp

rng = jax.random.PRNGKey(42)
A = jax.random.normal(rng, shape=(128, 16, 16))

# This is the function I would like to take the row-wise Jacobian of
def already_vmapped(x, A):
vmap_mul = jax.vmap(jnp.matmul)
return vmap_mul(A, x)

# This is how I am doing it now
function = lambda x, A: jax.vmap(jax.jacobian(lambda x, A: already_vmapped(x[None,], A[None,]),
argnums=0))(x, A).squeeze()

x = jnp.ones((128, 16))
print(jnp.allclose(function(x, A), A)) # True

# A slightly cleaner way that is too memory intensive
function_mem = lambda x, A: jnp.diagonal(jax.jacobian(already_vmapped, argnums=0)(x, A),
offset=0, axis1=0, axis2=2).transpose(2, 0, 1)

print(jnp.allclose(function_mem(x, A), A)) # True


I understand that the absolute cleanest way would just to not vmap the original function in the first place, but for whatever reason that is not easy to undo right now given my codebase. Any other suggestions are welcome!

Continue reading...