jaxray 0.1.0

A minimal JAX interpretability library

jaxray

A minimal JAX interpretability library for activation caching and patching.

Why jaxray?

Current JAX interpretability tooling has significant friction:

• Penzai requires learning new paradigms (pytrees as models, named axes) and has poor interop with existing Flax/Haiku models
• XLens is a direct port of TransformerLens with limited model coverage
• No JAX equivalent to PyTorch's NNsight/pyvene for wrapping arbitrary models

Meanwhile, PyTorch researchers have mature, frictionless options (TransformerLens, NNsight, pyvene).

The Core Problem

JAX's functional purity makes activation caching awkward. Side effects break jax.jit:

cache = {}
def model(params, x):
    x = layer(params, x)
    cache['layer'] = x  # side effect  breaks jax.jit
    return x

The Solution

Return the cache explicitly  it's pure and works with all JAX transforms:

def model(params, x):
    cache = {}
    x = layer(params, x)
    cache['layer'] = x
    return x, cache  # pure  jax.jit works

jaxray provides simple scaffolding around this pattern.

Installation

pip install jaxray

Usage

The @intercept Decorator

Decorate your model function and use _save to mark activation points:

from jaxray import intercept
import jax.numpy as jnp
import jax

@intercept
def mlp(params, x, _save):
    x = jnp.dot(x, params['w1']) + params['b1']
    x = _save('pre_relu', x)

    x = jax.nn.relu(x)
    x = _save('post_relu', x)

    x = jnp.dot(x, params['w2']) + params['b2']
    return x

# Basic forward pass  no caching
output, cache = mlp(params, x)  # cache is empty dict

# Cache specific activations
output, cache = mlp(params, x, cache={'pre_relu', 'post_relu'})
print(cache['pre_relu'].shape)

# Patch  replace an activation with a fixed value
output, cache = mlp(params, x, patch={'post_relu': jnp.zeros(hidden_dim)})

# Patch  transform an activation
output, cache = mlp(params, x, patch={'post_relu': lambda act: act * 0.5})

# Combine caching and patching
output, cache = mlp(params, x,
    cache={'pre_relu'},
    patch={'post_relu': lambda act: act + steering_vector}
)

With JAX Transforms

Use functools.partial to fix cache/patch before applying JAX transforms:

from functools import partial

# jit
mlp_cached = partial(mlp, cache={'pre_relu'})
jitted_mlp = jax.jit(mlp_cached)
output, cache = jitted_mlp(params, x)

# vmap
batched_mlp = jax.vmap(partial(mlp, cache={'pre_relu'}))
output, cache = batched_mlp(params, batch_x)

# grad (works directly)
def loss_fn(params, x):
    output, _ = mlp(params, x)
    return output.sum()

grads = jax.grad(loss_fn)(params, x)

Design Principles

1. No new paradigms  uses standard JAX/Flax patterns
2. Minimal API  a decorator and a few helpers, not a framework
3. Composable  works with jax.jit, jax.grad, jax.vmap
4. Toy model friendly  zero friction for simple MLPs and transformers

Current Status

Implemented:

• @intercept decorator with caching and patching
• Full compatibility with JAX transforms

Planned:

• patching_sweep helper for activation patching experiments
• steering_hook, zero_ablate, mean_ablate helpers
• Pre-built toy models (MLP, Transformer)