How to access the gradient value when using flax.nnx.value_and_grad

[adamhadani]

Hello,

I have recently started using the jax/flax/optax stack (stax?) for some deep reinforcement learning projects.

I ran into the following issue pretty early on. Notice I'm using the flax.nnx.value_and_grad "lifted" version of jax.value_and_grad that provides some boilerplate for state management and working with nnx.Modules. A single training step looks like this for example:

@nnx.jit
def train_step(model, optimizer, obs, actions, scaled_returns):
        grad_fn = nnx.value_and_grad(loss_fn, has_aux=True)
        (loss, log_probs), grads = grad_fn(model, obs, actions, scaled_returns)
        optimizer.update(grads)

        return loss, log_probs, grads

I would like to track the gradient norms (e.g. for tracking in tensorboard). how do I access the actual numeric value for the computed gradients? the grads value returned here seems to be a function/callable that is used by the optimizer.update rather than the actual computed gradient?

Sorry if this is a silly question or belongs in some different flax forum!

Cheers,

Adam

[cgarciae]

Hi @adamhadani, grads here is a nnx.State (possibly nested) mapping when contains the arrays representing the gradient.

from flax import nnx
import jax.numpy as jnp

class Foo(nnx.Module):
  def __init__(self):
    self.a = nnx.Param(jnp.ones((2, 2)))

def loss_fn(model):
  return jnp.sum(model.a.value)

model = Foo()
grads = nnx.grad(loss_fn)(model)

print(grads)

Output:

State({
  'a': VariableState( # 4 (16 B)
    type=Param,
    value=Array([[1., 1.],
           [1., 1.]], dtype=float32)
  )
})

[adamhadani]

Got it @cgarciae , many thanks!

In that case, does it make sense to try and calculate the overall gradient norm using something like the following?

...
grads = nnx.grad(loss_fn)(model)

grad_norm = jnp.sqrt(jax.tree.reduce(lambda x, y: (x**2.).sum() + (y**2.).sum(), grads))

[adamhadani]

If helps anybody (or for suggestions for improvements!), I ended up with the following snippet to compute the gradient norm across an arbitrary pytree gradient structure as returned from the nnx.grad or nnx.value_and_grad functions:

def calculate_gradient_norm_nnx(grads):
    return jnp.sqrt(
        jnp.array([
            jnp.linalg.norm(g) ** 2 
            if g is not None else 0.0 
            for g in jax.tree_util.tree_flatten(grads)[0]
        ]).sum()
    )