Getting the expected dimensions of the Jacobian with JAX?

Question

I am trying to get the Jacobian for a simple parameterization function within JAX. The code is as follows:

# imports 
import jax 
import jax.numpy as jnp
from jax import random

# simple parameterization function 
def reparameterize(v_params):
    theta = v_params[0] + jnp.exp(v_params[1]) * eps
    return theta

Suppose I initialize eps to be a vector of shape (3,) and v_params to be of shape (3, 2):

key = random.PRNGKey(2022)
eps = random.normal(key, shape=(3,))
key, _ = random.split(key)
v_params = random.normal(key, shape=(3, 2))

I want the Jacobian to be an array of shape (3, 2) but by using

jacobian(vmap(reparameterize))(v_params)

returns an array of shape (3, 3, 3, 2). If I re-initialize with only a single eps:

key, _ = random.split(key)
eps = random.normal(key, shape=(1, ))
key, _ = random.split(key)
v_params = random.normal(key, shape=(2, ))

and call jacobian(reparameterize)(v_params) I get what I want, e.g., an array of shape (2, ). Effectively looping over all eps and stacking the results of each Jacobian gives me the desired Jacobian (and shape). What am I missing here? Thanks for your help!

score 1 · Accepted Answer · answered Oct 01 '22 at 14:12

For a function f that maps an input of shape shape_in to an output of shape shape_out, the jacobian will have shape (*shape_out, *shape_in).

In your case, vmap(reparameterize) takes an array of shape (3, 2) and returns an array of shape (3, 3), so the output of the jacobian is an array of shape (3, 3, 3, 2).

It's hard to tell from your question what computation you were intending, but if you want a jacobian the same shape as the input, you need a function that maps the input to a scalar. Perhaps the sum is what you had in mind?

result = jacobian(lambda x: vmap(reparameterize)(x).sum())(v_params)
print(result.shape)
# (3, 2)

Getting the expected dimensions of the Jacobian with JAX?

1 Answers1