peex 0.0.2

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

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.