peex 0.0.3

Runtime hooks for Equinox modules.

Peex

Peex is a library for adding runtime evaluation hooks to Equinox modules, allowing you to peek inside!

This includes:

• forward and backward pass hooks
• runtime logging to inspect model activations/gradients
• simple wrapping interface; no edits to original modules
• compatible with jit, vmap, grad, etc

Installation

uv pip install peex

Example

Forward hooks
We can add runtime hooks and logging to the forward evaluation of an Equinox module. This edits/logs the model activations x as they pass forwards through the module.

import equinox as eqx
import jax

import peex as px


class MLP(eqx.Module):
    linear1: eqx.nn.Linear
    linear2: eqx.nn.Linear

    def __init__(self, in_features, width, out_features, *, key):
        key1, key2 = jax.random.split(key, 2)
        self.linear1 = eqx.nn.Linear(in_features, width, key=key1)
        self.linear2 = eqx.nn.Linear(width, out_features, key=key2)

    def __call__(self, x):
        x = self.linear1(x)
        x = jax.nn.relu(x)
        x = self.linear2(x)
        return x


mlp_key, x_key = jax.random.split(jax.random.key(0), 2)
mlp = MLP(1, 2, 1, key=mlp_key)
x = jax.random.normal(x_key, shape=(1,))

hooked_mlp = px.ForwardHook(mlp, hook_fn=px.IdentityHook(), log_fn=px.IdentityLog())
output, log = hooked_mlp(x)

# output = [0.23334768]
# log = {'linear1': Array([1.2571917, 1.2331474], dtype=float32), 
#        'linear2': Array([0.23334768], dtype=float32)}

Backward hooks
We can also add runtime hooks and logging to the backward pass of an Equinox module. This edits/logs the gradients (cotangent values of x) as they pass backwards through the module.

@eqx.filter_value_and_grad
def grad_loss(hooked_mlp, x):
    output = hooked_mlp(x)
    return jnp.mean(output)


hooked_mlp = px.BackwardHook(
    mlp,
    hook_fn=px.IdentityHook(),
    log_fn=px.IdentityLog(),
)
hooked_mlp = hooked_mlp.init(x)
loss, grads = grad_loss(hooked_mlp, x)

model_grads = grads.model
log_grads = grads.log

# log_grads = {'linear1': Array([-0.68400913,  0.4777369 ], dtype=float32), 
#              'linear2': Array([1.], dtype=float32)}

Note that for the backward hook we must run a post-init hooked_mlp = hooked_mlp.init(x). This runs a jax.eval_shape call to populate the log dictionary with the correct gradient shapes. The log dictionary can then be populated with cotangent values on the backward pass.

See the examples/ directory for more in-depth examples.

Writing custom hooks

The defining rule of a hook_fn is that the output shapes/dtypes exactly match their input shapes/dtypes. This can be understood by tracing the forward logic:

Pre-hook

# original module
x -> module(x)
# with hook
x -> module(hook_fn(x))

Here, x are tuple of arguments. So, hook_fn should map tuple -> tuple.

Post-hook

# original module
x -> module(x)
# with hook
x -> hook_fn(module(x))

Here, module(x) is the output of module with unconstrained type output_type. So, hook_fn should map output_type -> output_type. The same logic applies on the backward pass.

Writing custom logs

A log_fn can return an output with any shape/dtype. The only requirement is that it accepts the correct input shapes/dtypes. This depends on whether you are running a pre or post hook - see the discussion on custom hooks above.