self.sow in custom_vjp function unable to capture gradients

[runiteking1]

I want to save and analyze the derivatives in the backward pass of a Jacobian for a simple NN (akin to ##5336). I can easily print out the values I want using a custom_vjp wrapper around each layer. However, saving it using self.sow somehow is not working. My current workaround is to save the values I want in a global list, but this seems extremely un-Pythonic. Below is a MWE where the self.sow on line 47 doesn't seem to capture anything.

If anybody know how to actually use sow (or something else that's not just a global list), I greatly appreciate it.

import jax
import jax.numpy as jnp
import flax.linen as nn

# Workaround is to use a global list
adjoint_lists = []

class AdjointHook(nn.Module):
    layer: nn.Module

    @nn.compact
    def __call__(self, *args, **kwargs):
        # Define the function to which we'll attach the custom VJP
        @jax.custom_vjp
        def layer_with_hook(params, *args, **kwargs):
            return self.layer.apply(params, *args, **kwargs)

        # --- Define the Forward and Backward Passes ---
        def layer_fwd(params, *args, **kwargs):
            # Execute the original layer's forward pass
            output = self.layer.apply(params, *args, **kwargs)

            # Return output and residuals for the backward pass
            return output, (params, args, kwargs)

        def layer_bwd(res, g):
            # res contains residuals: (params, args, kwargs)
            # g is the incoming gradient (the adjoint we want to capture)
            params, args, kwargs = res

            # Calculate the VJP of the original wrapped layer
            # This computes the gradients w.r.t. params and inputs
            _, vjp_fun = jax.vjp(
                lambda p, *a, **kw: self.layer.apply(p, *a, **kw), params, *args, **kwargs
            )

            # The VJP function returns a tuple of gradients
            grad_params, *grad_args = vjp_fun(g)

            print(f"--- (Partially) Captured Adjoint for layer: {self.layer.name} ---")
            print(f'{g.val.flatten()[0]}, {grad_args[0][0].val.flatten()[0]}')
            print("-" * 30)

            # Attempt to save to collection adjoints
            # ------------- DOES NOTHING? -------------------------------
            # Trying to capture something...
            self.sow('adjoints', 'captured_adjoints', g.val.flatten()[0])

            # Work around
            adjoint_lists.append(grad_args[0])

            return (grad_params,) + tuple(grad_args)

        # Attach the custom forward and backward functions
        layer_with_hook.defvjp(layer_fwd, layer_bwd)

        # Get the parameters for the wrapped layer
        layer_params = self.param('wrapped_layer', self.layer.init, *args, **kwargs)

        return layer_with_hook(layer_params, *args, **kwargs)


class MLPLayer(nn.Module):
    num_units: int

    @nn.compact
    def __call__(self, x):
        f = nn.Dense(self.num_units)(x)
        f = nn.tanh(f)
        f = nn.Dense(self.num_units)(f)
        return f

class SimpleMLP(nn.Module):
    num_layers: int
    num_units: int

    def setup(self):
        # Create a list of Dense layers
        self.layers = [
            AdjointHook(MLPLayer(self.num_units, name=f"layer_{i:02})")) for i in range(self.num_layers)
        ]

    def __call__(self, x):
        # Pass the input through each layer
        for i, layer in enumerate(self.layers):
            x = layer(x)

            # Store intermediate outputs of layer; this works well
            self.sow('intermediates', f'layer_{i + 1}_output', x)
        return x

if __name__ == '__main__':
    L, n_samples, n_dim = 4, 2, 5
    X = jnp.linspace(0, 1, n_samples* n_dim).reshape((n_samples, n_dim))

    key = jax.random.PRNGKey(0)
    model = SimpleMLP(num_layers=L, num_units=n_dim)
    params = model.init(key, jnp.ones((1, n_dim)))

    # We can find out the layer outputs of inference using sow
    # predictions, intermediates = model.apply(params, X, mutable=['intermediates'])
    # print(intermediates)

    # Simple code to find intermediate of backward using jacobian
    jacobian = jax.jacobian(model.apply, argnums=0, has_aux=True)(params, X, mutable=['adjoints'])
    print(f'Using the sow does not record anything (?): {jacobian[1]}')
    print(f'Our global list captures it: {[float(adj.val.flatten()[0]) for adj in adjoint_lists]}')

The outputs on my end are

--- (Partially) Captured Adjoint for layer: layer_03) ---
1.0, -0.31861528754234314
------------------------------
--- (Partially) Captured Adjoint for layer: layer_02) ---
-0.31861528754234314, 0.10002732276916504
------------------------------
--- (Partially) Captured Adjoint for layer: layer_01) ---
0.10002732276916504, -0.07864561676979065
------------------------------
--- (Partially) Captured Adjoint for layer: layer_00) ---
-0.07864561676979065, -0.09821531921625137
------------------------------
Using the sow does not record anything (?): {}
Our global list captures it: [-0.31861528754234314, 0.10002732276916504, -0.07864561676979065, -0.09821531921625137]