Inspection and custom layers

Implement a differentiable custom layer, register non-trainable state, inspect its hierarchy, and use the pure state boundary.

These outputs were saved from local execution. The documentation renderer does not execute notebook cells.

Show code cell source

Hide code cell source

from pathlib import Path
import sys
ROOT = Path.cwd() if (Path.cwd() / 'examples').exists() else Path.cwd().parent
sys.path.insert(0, str(ROOT))
import jax
import jax.numpy as jnp
import numpy as np
import jaxsaw as js
from jaxsaw import nn, optim, data, checkpoint, inspection
from jaxsaw.nn import functional as F
print(f"JAXsaw {js.__version__} · JAX {jax.__version__}")
print(jax.devices())

JAXsaw 0.1.0 · JAX 0.5.3
[CudaDevice(id=0)]

Implement an affine layer

The layer uses ordinary JAX operations. Parameters receive gradients automatically; no manual backward method is needed.

class Affine(nn.Module):
    def __init__(self,features,*,rngs):
        super().__init__()
        self.features = features
        self.scale = nn.Parameter(jnp.ones(features))
        self.offset = nn.Parameter(jnp.zeros(features))
    def forward(self,x):
        return x*self.scale.value + self.offset.value
model = nn.Sequential(Affine(3,rngs=js.Rngs(params=0)),
                      nn.Linear(3,1,rngs=js.Rngs(params=1)))
print(inspection.summary(model))

Module | Type | Own parameters
items.0 | Affine | 6
items.1 | Linear | 4
<root> | Sequential | 0
Total unique parameters: 10

Inspect without mutating the live model

output_spec runs an evaluation clone. The architecture description records ownership, shapes, configuration, and aliases; it is not a full operation-level execution graph.

print(inspection.output_spec(model,jnp.ones((4,3))))
import json
print(json.dumps(inspection.describe(model),indent=2))

{'shape': [4, 1], 'dtype': 'float32'}
{
  "schema_version": 1,
  "scope": "module hierarchy; not an execution graph",
  "root": {
    "type": "jaxsaw.nn.module.Sequential",
    "fields": {
      "items": [
        {
          "type": "__main__.Affine",
          "fields": {
            "features": 3,
            "scale": {
              "kind": "Param",
              "shape": [
                3
              ],
              "dtype": "float32"
            },
            "offset": {
              "kind": "Param",
              "shape": [
                3
              ],
              "dtype": "float32"
            }
          }
        },
        {
          "type": "jaxsaw.nn.layers.Linear",
          "fields": {
            "in_features": 3,
            "out_features": 1,
            "weight": {
              "kind": "Param",
              "shape": [
                1,
                3
              ],
              "dtype": "float32"
            },
            "bias": {
              "kind": "Param",
              "shape": [
                1
              ],
              "dtype": "float32"
            }
          }
        }
      ]
    }
  }
}

Differentiate custom parameters

Only Parameter leaves are differentiated by default. Buffers such as running statistics are persistent state but not optimizer parameters.

loss,grads = js.value_and_grad(lambda m:jnp.square(m(jnp.ones((4,3)))).mean())(model)
optimizer = optim.SGD(model.parameters(),lr=.01)
optimizer.step(grads)
print('Loss:',float(loss))
print('Learned affine scale:',np.asarray(model[0].scale.value))

Loss: 0.08691398799419403
Learned affine scale: [0.9986205 1.0018053 0.9978359]

Cross the pure functional boundary

split separates static structure from dynamic state. apply returns output and updated state without mutating the original model. This is useful when integrating pure JAX computations.

graph,state = js.split(model)
x = jnp.ones((2,3))
output,next_state = js.apply(graph,state,x)
np.testing.assert_allclose(output,model(x))
print('Pure output:',np.asarray(output))
js.update(model,next_state)

Pure output: [[0.2679077]
 [0.2679077]]

Register persistent buffers

Use buffers for state you intend to checkpoint but not optimize. Side effects outside registered state, such as appending to a Python log, do not belong inside compiled steps.

class Center(nn.Module):
    def __init__(self,mean):
        super().__init__()
        self.mean = nn.Buffer(jnp.asarray(mean,dtype=jnp.float32))
    def forward(self,x):
        return x-self.mean.value
center = Center([1.,2.,3.])
print('Buffers:',[(name,value.value.shape) for name,value in center.named_buffers()])
print('Trainable arrays:',len(list(center.parameters())))
print(center(jnp.array([[2.,4.,6.]])))

Buffers: [('mean', (3,))]
Trainable arrays: 0
[[1. 2. 3.]]