# 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`:

```python
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:

```python
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

```bash
pip install jaxray
```

## Usage

### The `@intercept` Decorator

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

```python
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:

```python
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)
