Can we use flax.linen.vjp to calculate the HVP of a loss wrt a network's weights?

[newalexander]

The jax documentation has a nice example of calculating the Hessian-vector product of a function using jax.grad and jax.vjp. I see that flax also has a vjp function; however, its functionality and use are not completely clear to me, such as the need for tangents, variable_tangents, and variables all as arguments.

Is it possible to use flax.linen.vjp and jax.grad to calculate the Hessian-vector product of a neural network's loss function with respect to its weights and another vector the same size as the weights? A nominal example below:

import jax
import flax

def hvp(loss, primals, tangents):
    return flax.linen.jvp(jax.grad(loss), primals, tangents)[1]

where primals and tangents are both FrozenDicts with the same set of keys and value shapes, and loss is a a function that accepts primals and calls an nn.Module.

Thank you!

[jheek]

The linen versions of jax transforms are only required when you want to use them inside a linen Module. In this case you should simply use jax.jvp and jax.grad

[newalexander]

@jheek Thank you for your response! Unfortunately, jax.jvp doesn't seem to quite work still, or I'm misunderstanding your suggestion.

In the below code, I get the error TypeError: primal and tangent arguments to jax.jvp must be tuples or lists; found FrozenDict and FrozenDict. when trying to call hvp.

import numpy as np

from jax import random, grad, jvp, numpy as jnp
from flax import linen as nn

np.random.seed(1234)

n_input, n_output, n_batch = 3, 5, 11
points = jnp.array(np.random.randn(n_batch, n_input))

key = random.PRNGKey(0)
model = nn.Dense(n_output)
params = model.init(key, points)


def mean(ps):
    return model.apply(ps, points).mean()

def hvp(f, primals, tangents):
    return jvp(grad(f), primals, tangents)

hvp(mean, params, params)

[jheek]

That's because you need to pass in something with the same data structure as *args which always is a tuple. So if you have f(params) then primals = (params,)

[newalexander]

@jheek Thank you! This now works and makes perfect sense.