Compute gradients wrt inputs instead of params

[andsteing]

(question originally asked by Bogdan Mazoure)

Hi, if I have a network, e.g. y=nn.Dense(256)(x), what is the most straightforward way to compute the gradient \grad_x of y as opposed to wrt the parameters of the nn.Dense? E.g. this arises when using Fisher divergence

[andsteing]

(answer originally provided by @levskaya)

setup code

import jax
from jax import lax, numpy as jnp, random
import flax
from flax import linen as nn

defining a layer_fn(x, variables)

layer = nn.Dense(32)

k0, k1, k2, k3 = random.split(random.PRNGKey(0), 4)
x = random.uniform(k0, (2,16))
variables = layer.init(k1, x)
zeros_like_variables = jax.tree_map(jnp.zeros_like, variables)

layer_fn = lambda x, variables: layer.apply(variables, x)

jvp

We have tangents we'd like to push forward.

tangent = random.uniform(k2, (2, 16))   # or "dx" to map forwards

y, tangents_out = jax.jvp(layer_fn,
                          (x, variables),                  # inputs
                          (tangent, zeros_like_variables)  # tangents_in
                          )
tangents_out.shape

(2, 32)

we can also break apart the above and call linearize to get the linearization of a function

y, linear_fn = jax.linearize(layer_fn, x, variables)
tangents_out = linear_fn(tangent, zeros_like_variables)
tangents_out.shape

(2, 32)

vjp

We have cotangents we'd like to pull-back.

jax.grad is a special case of jax.vjp used for functions with scalar outputs.

cotangent = random.uniform(k3, (2, 32)) # or "dy" to map back

output, vjp_fun = jax.vjp(layer_fn, x, variables)
dx, dvariables = vjp_fun(cotangent)
dx.shape

(2, 16)

jacobian

We'd like to get the actual jacobian of dy / dx for a given x, variables.

(You can also check out docs for jax.jacfwd and jax.jacrev.)

jac_dy_dx = jax.jacobian(layer_fn, 
                         argnums=(0,)  # jacobian w.r.t. "x"
                         )(x, variables)[0]
jac_dy_dx.shape

(2, 32, 2, 16)