Surrogate Gradient Tutorial¶

This tutorial trains a small spiking neural network on the Spiking Heidelberg Digits (SHD) dataset using backpropagation through time and surrogate gradients. The network is built with the new Spyx Flax-NNX API and the dataset is streamed via spyx.data.SHD_loader, a thin wrapper around Google Grain + Tonic.

Required extras:

pip install "spyx[loaders]"

In [ ]:

# Colab / fresh environment: install Spyx + this tutorial's extra deps.
# `scikit-learn` is the correct PyPI name (imported as `sklearn`).
%pip install -q "spyx[loaders]" scikit-learn matplotlib

# Colab / fresh environment: install Spyx + this tutorial's extra deps. # `scikit-learn` is the correct PyPI name (imported as `sklearn`). %pip install -q "spyx[loaders]" scikit-learn matplotlib

In [1]:

import jax
import jax.numpy as jnp
import numpy as np
import optax
from flax import nnx
from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
from tqdm import trange
import matplotlib.pyplot as plt

import spyx
import spyx.nn as snn

import jax import jax.numpy as jnp import numpy as np import optax from flax import nnx from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix from tqdm import trange import matplotlib.pyplot as plt import spyx import spyx.nn as snn

/home/kade/Code/spyx/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Inspect a surrogate gradient¶

Spyx exposes a family of surrogate-gradient activations in spyx.axn. Each is built from jax.custom_gradient, so JAX can trace it transparently and JIT-compile the whole training loop.

In [2]:

from jax import make_jaxpr

example = jnp.linspace(-1, 1, 5)
print(make_jaxpr(spyx.axn.superspike())(example))

from jax import make_jaxpr example = jnp.linspace(-1, 1, 5) print(make_jaxpr(spyx.axn.superspike())(example))

{ lambda ; a:f32[5]. let
    b:f32[5] = jit[
      name=wrapped_fun
      jaxpr={ lambda ; a:f32[5]. let
          b:f32[5] = custom_vjp_call[
            name=wrapped_fun
            bwd=bwd
            call_jaxpr={ lambda ; c:f32[5]. let
                d:bool[5] = gt c 0.0:f32[]
                e:i32[5] = jit[
                  name=_where
                  jaxpr={ lambda ; d:bool[5] f:i32[] g:i32[]. let
                      h:i32[5] = broadcast_in_dim f
                      i:i32[5] = broadcast_in_dim g
                      e:i32[5] = select_n d i h
                    in (e,) }
                ] d 1:i32[] 0:i32[]
                j:f32[5] = convert_element_type[
                  new_dtype=float32
                  weak_type=False
                ] e
              in (j,) }
            fwd=fwd
            symbolic_zeros=False
          ] a
        in (b,) }
    ] a
  in (b,) }

Data loading¶

spyx.data.SHD_loader returns an object whose train_epoch() / test_epoch() methods yield State(obs, labels) per batch. To minimise host→device transfer the spike tensor is bit-packed along the time axis with numpy.packbits; we recover dense binary spikes with jnp.unpackbits inside the forward pass.

In [3]:

BATCH = 256
SAMPLE_T = 128
CHANNELS = 128

shd_dl = spyx.data.SHD_loader(batch_size=BATCH, sample_T=SAMPLE_T, channels=CHANNELS)
print("observation shape:", shd_dl.obs_shape, "num classes:", int(shd_dl.act_shape[0]))

BATCH = 256 SAMPLE_T = 128 CHANNELS = 128 shd_dl = spyx.data.SHD_loader(batch_size=BATCH, sample_T=SAMPLE_T, channels=CHANNELS) print("observation shape:", shd_dl.obs_shape, "num classes:", int(shd_dl.act_shape[0]))

observation shape: (128,) num classes: 20

In [4]:

# Prestage the whole dataset onto the accelerator once (the "Prestage" section
# below explains why). We reuse these arrays both for this visualisation and for
# the training loop. SHD is tiny, so a one-shot bulk load is quick.
train_obs, train_labels = shd_dl.prestage("train")
test_obs, test_labels = shd_dl.prestage("test")
print("packed train obs:", train_obs.shape, train_obs.dtype)

# Find a sample with actual spikes. A handful of SHD recordings contain
# NaN / inf entries that spyx.data drops, so the first row in a batch may be
# silent. Scan the first prestaged batch until we find one with events.
first_batch = np.asarray(train_obs[0])
idx = next(
    (i for i in range(first_batch.shape[0]) if first_batch[i].any()),
    0,
)
print(f"visualising sample index {idx}, label={int(train_labels[0][idx])}")

# Unpack the packed time axis and slice to SAMPLE_T timesteps.
spikes_one = jnp.unpackbits(train_obs[0][idx], axis=0)[:SAMPLE_T]
plt.imshow(spikes_one.T, aspect="auto", cmap="Greys")
plt.xlabel("time")
plt.ylabel("channel")
plt.title(f"single SHD sample (batch index {idx})")
plt.show()

# Prestage the whole dataset onto the accelerator once (the "Prestage" section # below explains why). We reuse these arrays both for this visualisation and for # the training loop. SHD is tiny, so a one-shot bulk load is quick. train_obs, train_labels = shd_dl.prestage("train") test_obs, test_labels = shd_dl.prestage("test") print("packed train obs:", train_obs.shape, train_obs.dtype) # Find a sample with actual spikes. A handful of SHD recordings contain # NaN / inf entries that spyx.data drops, so the first row in a batch may be # silent. Scan the first prestaged batch until we find one with events. first_batch = np.asarray(train_obs[0]) idx = next( (i for i in range(first_batch.shape[0]) if first_batch[i].any()), 0, ) print(f"visualising sample index {idx}, label={int(train_labels[0][idx])}") # Unpack the packed time axis and slice to SAMPLE_T timesteps. spikes_one = jnp.unpackbits(train_obs[0][idx], axis=0)[:SAMPLE_T] plt.imshow(spikes_one.T, aspect="auto", cmap="Greys") plt.xlabel("time") plt.ylabel("channel") plt.title(f"single SHD sample (batch index {idx})") plt.show()

packed train obs: (31, 256, 16, 128) uint8
visualising sample index 0, label=11

SNN model¶

Two LIF hidden layers and a leaky-integrate (LI) readout. spyx.nn.Sequential threads per-layer hidden state through __call__, and spyx.nn.run scans over the time axis with jax.lax.scan (it expects a [T, B, ...] input). NNX layers are eagerly initialised, so the first Linear needs concrete in_features.

In [5]:

class SHDSNN(nnx.Module):
    def __init__(self, in_dim, hidden, n_classes, *, rngs):
        self.core = snn.Sequential(
            nnx.Linear(in_dim, hidden, use_bias=False, rngs=rngs),
            snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs),
            nnx.Linear(hidden, hidden, use_bias=False, rngs=rngs),
            snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs),
            nnx.Linear(hidden, n_classes, use_bias=False, rngs=rngs),
            snn.LI((n_classes,), rngs=rngs),
        )

    def __call__(self, x_BTC):
        """x_BTC has shape (batch, time, channels). Returns traces (batch, time, classes)."""
        x_TBC = jnp.transpose(x_BTC, (1, 0, 2))
        traces, _ = snn.run(self.core, x_TBC)
        return jnp.transpose(traces, (1, 0, 2))


model = SHDSNN(in_dim=CHANNELS, hidden=64, n_classes=int(shd_dl.act_shape[0]), rngs=nnx.Rngs(0))

class SHDSNN(nnx.Module): def __init__(self, in_dim, hidden, n_classes, *, rngs): self.core = snn.Sequential( nnx.Linear(in_dim, hidden, use_bias=False, rngs=rngs), snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs), nnx.Linear(hidden, hidden, use_bias=False, rngs=rngs), snn.LIF((hidden,), activation=spyx.axn.triangular(), rngs=rngs), nnx.Linear(hidden, n_classes, use_bias=False, rngs=rngs), snn.LI((n_classes,), rngs=rngs), ) def __call__(self, x_BTC): """x_BTC has shape (batch, time, channels). Returns traces (batch, time, classes).""" x_TBC = jnp.transpose(x_BTC, (1, 0, 2)) traces, _ = snn.run(self.core, x_TBC) return jnp.transpose(traces, (1, 0, 2)) model = SHDSNN(in_dim=CHANNELS, hidden=64, n_classes=int(shd_dl.act_shape[0]), rngs=nnx.Rngs(0))

Prestage the dataset on device¶

The Spyx paper's throughput story rests on two tricks: (1) load the whole dataset into GPU memory once, (2) scan over it inside a JIT-compiled kernel per epoch. SHD is only ~130 MB packed, so it fits comfortably on any training accelerator. SHD_loader.prestage(split) does the one-shot bulk load via PyTorch's DataLoader (single call, batch_size=len(split), tonic's native PadTensors collator), which is much faster than iterating the grain streaming pipeline for this use case.

In [6]:

import time

t0 = time.perf_counter()
print("Prestaging SHD train split...")
train_obs, train_labels = shd_dl.prestage("train")
print(f"  train done in {time.perf_counter() - t0:.1f}s")

t0 = time.perf_counter()
print("Prestaging SHD test split...")
test_obs, test_labels = shd_dl.prestage("test")
print(f"  test  done in {time.perf_counter() - t0:.1f}s")

print(f"train obs: {train_obs.shape} {train_obs.dtype}")
print(f"test  obs: {test_obs.shape} {test_obs.dtype}")
print(f"~{train_obs.nbytes / 1e6:.1f} MB train + {test_obs.nbytes / 1e6:.1f} MB test on device")

import time t0 = time.perf_counter() print("Prestaging SHD train split...") train_obs, train_labels = shd_dl.prestage("train") print(f" train done in {time.perf_counter() - t0:.1f}s") t0 = time.perf_counter() print("Prestaging SHD test split...") test_obs, test_labels = shd_dl.prestage("test") print(f" test done in {time.perf_counter() - t0:.1f}s") print(f"train obs: {train_obs.shape} {train_obs.dtype}") print(f"test obs: {test_obs.shape} {test_obs.dtype}") print(f"~{train_obs.nbytes / 1e6:.1f} MB train + {test_obs.nbytes / 1e6:.1f} MB test on device")

Prestaging SHD train split...

  train done in 12.9s
Prestaging SHD test split...

  test  done in 3.1s
train obs: (31, 256, 16, 128) uint8
test  obs: (8, 256, 16, 128) uint8
~16.3 MB train + 4.2 MB test on device

Training step + eval on prestaged data¶

Two small JIT-compiled helpers: one gradient step on a prestaged batch, one eval pass over the full test split. The key performance trick is that train_obs / test_obs already live on the accelerator, so the tonic + grain decode cost is paid once up front rather than per batch.

Expected result: with the default EPOCHS = 30, this 64-unit two-layer LIF network reaches roughly 60% test accuracy on SHD (a 20-way task; chance is 5%). SHD keeps improving to ~80% around 80 epochs — bump EPOCHS for a stronger model. Accuracy climbs slowly for the first ~5 epochs while the surrogate-gradient signal warms up, then rises steadily.

In [7]:

Loss = spyx.fn.integral_crossentropy()
Acc = spyx.fn.integral_accuracy()
optimizer = nnx.Optimizer(model, optax.lion(3e-4), wrt=nnx.Param)


def _unpack(packed_obs):
    """(B, T_packed, C) uint8 -> (B, T, C) float32."""
    return jnp.unpackbits(packed_obs, axis=1)[:, :SAMPLE_T, :].astype(jnp.float32)


@nnx.jit
def train_step(model, optimizer, packed_obs, targets):
    """One forward + backward + optimizer step on a prestaged batch."""
    def loss_fn(m):
        return Loss(m(_unpack(packed_obs)), targets)
    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss


@nnx.jit
def eval_epoch(model, obs_NBTC, labels_NB):
    """Mean accuracy / loss across every test batch, in one JIT kernel."""
    def per_batch(i):
        packed_obs = obs_NBTC[i]
        targets = labels_NB[i]
        traces = model(_unpack(packed_obs))
        acc, _ = Acc(traces, targets)
        return acc, Loss(traces, targets)

    stats = jax.vmap(per_batch)(jnp.arange(obs_NBTC.shape[0]))
    return stats[0].mean(), stats[1].mean()

Loss = spyx.fn.integral_crossentropy() Acc = spyx.fn.integral_accuracy() optimizer = nnx.Optimizer(model, optax.lion(3e-4), wrt=nnx.Param) def _unpack(packed_obs): """(B, T_packed, C) uint8 -> (B, T, C) float32.""" return jnp.unpackbits(packed_obs, axis=1)[:, :SAMPLE_T, :].astype(jnp.float32) @nnx.jit def train_step(model, optimizer, packed_obs, targets): """One forward + backward + optimizer step on a prestaged batch.""" def loss_fn(m): return Loss(m(_unpack(packed_obs)), targets) loss, grads = nnx.value_and_grad(loss_fn)(model) optimizer.update(model, grads) return loss @nnx.jit def eval_epoch(model, obs_NBTC, labels_NB): """Mean accuracy / loss across every test batch, in one JIT kernel.""" def per_batch(i): packed_obs = obs_NBTC[i] targets = labels_NB[i] traces = model(_unpack(packed_obs)) acc, _ = Acc(traces, targets) return acc, Loss(traces, targets) stats = jax.vmap(per_batch)(jnp.arange(obs_NBTC.shape[0])) return stats[0].mean(), stats[1].mean()

In [8]:

EPOCHS = 30  # SHD typically plateaus around 80 epochs; bump for best accuracy.
N_BATCHES = train_obs.shape[0]

history = []
key = jax.random.PRNGKey(0)
for epoch in trange(EPOCHS):
    key, shuffle_key = jax.random.split(key)
    perm = jax.random.permutation(shuffle_key, N_BATCHES)
    epoch_losses = []
    for idx in perm:
        epoch_losses.append(train_step(model, optimizer, train_obs[idx], train_labels[idx]))

    eval_acc, eval_loss = eval_epoch(model, test_obs, test_labels)
    history.append((
        float(jnp.mean(jnp.stack(epoch_losses))),
        float(eval_acc),
        float(eval_loss),
    ))

print("final:", history[-1])

EPOCHS = 30 # SHD typically plateaus around 80 epochs; bump for best accuracy. N_BATCHES = train_obs.shape[0] history = [] key = jax.random.PRNGKey(0) for epoch in trange(EPOCHS): key, shuffle_key = jax.random.split(key) perm = jax.random.permutation(shuffle_key, N_BATCHES) epoch_losses = [] for idx in perm: epoch_losses.append(train_step(model, optimizer, train_obs[idx], train_labels[idx])) eval_acc, eval_loss = eval_epoch(model, test_obs, test_labels) history.append(( float(jnp.mean(jnp.stack(epoch_losses))), float(eval_acc), float(eval_loss), )) print("final:", history[-1])

  0%|          | 0/30 [00:00<?, ?it/s]

  3%|▎         | 1/30 [00:01<00:57,  1.99s/it]

  7%|▋         | 2/30 [00:02<00:38,  1.38s/it]

 10%|█         | 3/30 [00:03<00:31,  1.16s/it]

 13%|█▎        | 4/30 [00:04<00:27,  1.05s/it]

 17%|█▋        | 5/30 [00:05<00:24,  1.01it/s]

 20%|██        | 6/30 [00:06<00:23,  1.04it/s]

 23%|██▎       | 7/30 [00:07<00:21,  1.06it/s]

 27%|██▋       | 8/30 [00:08<00:19,  1.12it/s]

 30%|███       | 9/30 [00:09<00:19,  1.10it/s]

 33%|███▎      | 10/30 [00:10<00:18,  1.11it/s]

 37%|███▋      | 11/30 [00:10<00:16,  1.14it/s]

 40%|████      | 12/30 [00:11<00:15,  1.14it/s]

 43%|████▎     | 13/30 [00:12<00:14,  1.15it/s]

 47%|████▋     | 14/30 [00:13<00:14,  1.14it/s]

 50%|█████     | 15/30 [00:14<00:13,  1.14it/s]

 53%|█████▎    | 16/30 [00:15<00:12,  1.16it/s]

 57%|█████▋    | 17/30 [00:16<00:11,  1.16it/s]

 60%|██████    | 18/30 [00:16<00:10,  1.14it/s]

 63%|██████▎   | 19/30 [00:17<00:09,  1.15it/s]

 67%|██████▋   | 20/30 [00:18<00:08,  1.18it/s]

 70%|███████   | 21/30 [00:19<00:07,  1.16it/s]

 73%|███████▎  | 22/30 [00:20<00:07,  1.12it/s]

 77%|███████▋  | 23/30 [00:21<00:06,  1.09it/s]

 80%|████████  | 24/30 [00:22<00:05,  1.09it/s]

 83%|████████▎ | 25/30 [00:23<00:04,  1.08it/s]

 87%|████████▋ | 26/30 [00:24<00:03,  1.10it/s]

 90%|█████████ | 27/30 [00:25<00:02,  1.11it/s]

 93%|█████████▎| 28/30 [00:25<00:01,  1.10it/s]

 97%|█████████▋| 29/30 [00:26<00:00,  1.10it/s]

100%|██████████| 30/30 [00:27<00:00,  1.10it/s]

100%|██████████| 30/30 [00:27<00:00,  1.08it/s]

final: (2.08160138130188, 0.5595703125, 2.2989468574523926)

In [9]:

plt.plot(history, label=["train loss", "val acc", "val loss"])
plt.title("SHD surrogate gradient training")
plt.legend()
plt.show()

plt.plot(history, label=["train loss", "val acc", "val loss"]) plt.title("SHD surrogate gradient training") plt.legend() plt.show()

Evaluation¶

Run the trained model on the test set, accumulate predictions, and visualise a confusion matrix.

In [10]:

# Prestaged, JIT-compiled inference over every test batch in one kernel.
@nnx.jit
def test_all(model, obs_NBTC, labels_NB):
    @nnx.scan(in_axes=(nnx.Carry, 0, 0), out_axes=(nnx.Carry, 0))
    def body(carry, packed_obs, targets):
        traces = model(_unpack(packed_obs))
        acc, preds = Acc(traces, targets)
        loss = Loss(traces, targets)
        return carry, (acc, loss, preds, targets)
    _, (accs, losses, all_preds, all_tgts) = body(None, obs_NBTC, labels_NB)
    return accs, losses, all_preds.reshape(-1), all_tgts.reshape(-1)


accs, losses, preds, tgts = test_all(model, test_obs, test_labels)
print("Accuracy:", float(accs.mean()), "Loss:", float(losses.mean()))

cm = confusion_matrix(np.asarray(tgts), np.asarray(preds))
ConfusionMatrixDisplay(cm).plot()
plt.show()

# Prestaged, JIT-compiled inference over every test batch in one kernel. @nnx.jit def test_all(model, obs_NBTC, labels_NB): @nnx.scan(in_axes=(nnx.Carry, 0, 0), out_axes=(nnx.Carry, 0)) def body(carry, packed_obs, targets): traces = model(_unpack(packed_obs)) acc, preds = Acc(traces, targets) loss = Loss(traces, targets) return carry, (acc, loss, preds, targets) _, (accs, losses, all_preds, all_tgts) = body(None, obs_NBTC, labels_NB) return accs, losses, all_preds.reshape(-1), all_tgts.reshape(-1) accs, losses, preds, tgts = test_all(model, test_obs, test_labels) print("Accuracy:", float(accs.mean()), "Loss:", float(losses.mean())) cm = confusion_matrix(np.asarray(tgts), np.asarray(preds)) ConfusionMatrixDisplay(cm).plot() plt.show()

Accuracy: 0.5595703125 Loss: 2.2989468574523926

Bonus: quantization-aware training¶

Once the fp32 model has converged, we can explore how it behaves under int8 quantization. spyx.quant.quantize wraps Google's qwix library with Spyx-aware defaults: every nnx.Linear / nnx.Conv layer becomes int8 weights + int8 activations, while the spiking dynamics (LIF, LI) stay in fp32 because their state recurrences don't survive integer rounding.

Install the optional extra first:

pip install "spyx[quant]"   # or: pip install "git+https://github.com/google/qwix"

The cells below:

1. Take the trained model from above and feed it through spyx.quant.quantize in QAT mode (quantization-aware training).
2. Fine-tune for a few epochs so the weights adapt to the int8 grid.
3. Evaluate on the test set to confirm accuracy holds up.
4. Repeat the experiment with the BitNet b1.58 ternary recipe (int2 weights via bitnet_ternary_rules()) for a memory-bound operating point.

All of this reuses the same train_step / eval_step logic — the quantized model is a drop-in nnx.Module.

In [11]:

import spyx

assert spyx.quant.available(), (
    "qwix not installed. Run `pip install \"spyx[quant]\"` first."
)

# Sample input for qwix tracing — pull one prestaged batch.
sample_events = _unpack(train_obs[0])

qmodel = spyx.quant.quantize(model, sample_events)
print("quantized model:", type(qmodel).__name__)

import spyx assert spyx.quant.available(), ( "qwix not installed. Run `pip install \"spyx[quant]\"` first." ) # Sample input for qwix tracing — pull one prestaged batch. sample_events = _unpack(train_obs[0]) qmodel = spyx.quant.quantize(model, sample_events) print("quantized model:", type(qmodel).__name__)

quantized model: SHDSNN

In [12]:

# Fine-tune the quantized model, reusing the prestaged data and JIT'd step
# helper. A handful of epochs is usually enough for int8 to recover most of
# the fp32 accuracy.
qat_optimizer = nnx.Optimizer(qmodel, optax.lion(1e-4), wrt=nnx.Param)

QAT_EPOCHS = 5
qat_history = []
for epoch in trange(QAT_EPOCHS, desc="QAT"):
    key, shuffle_key = jax.random.split(key)
    perm = jax.random.permutation(shuffle_key, N_BATCHES)
    epoch_losses = []
    for idx in perm:
        epoch_losses.append(train_step(qmodel, qat_optimizer, train_obs[idx], train_labels[idx]))
    eval_acc, eval_loss = eval_epoch(qmodel, test_obs, test_labels)
    qat_history.append((
        float(jnp.mean(jnp.stack(epoch_losses))),
        float(eval_acc),
        float(eval_loss),
    ))

print("int8 QAT final:", qat_history[-1])

# Fine-tune the quantized model, reusing the prestaged data and JIT'd step # helper. A handful of epochs is usually enough for int8 to recover most of # the fp32 accuracy. qat_optimizer = nnx.Optimizer(qmodel, optax.lion(1e-4), wrt=nnx.Param) QAT_EPOCHS = 5 qat_history = [] for epoch in trange(QAT_EPOCHS, desc="QAT"): key, shuffle_key = jax.random.split(key) perm = jax.random.permutation(shuffle_key, N_BATCHES) epoch_losses = [] for idx in perm: epoch_losses.append(train_step(qmodel, qat_optimizer, train_obs[idx], train_labels[idx])) eval_acc, eval_loss = eval_epoch(qmodel, test_obs, test_labels) qat_history.append(( float(jnp.mean(jnp.stack(epoch_losses))), float(eval_acc), float(eval_loss), )) print("int8 QAT final:", qat_history[-1])

QAT:   0%|          | 0/5 [00:00<?, ?it/s]

QAT:  20%|██        | 1/5 [00:01<00:05,  1.34s/it]

QAT:  40%|████      | 2/5 [00:02<00:03,  1.08s/it]

QAT:  60%|██████    | 3/5 [00:03<00:01,  1.01it/s]

QAT:  80%|████████  | 4/5 [00:03<00:00,  1.06it/s]

QAT: 100%|██████████| 5/5 [00:04<00:00,  1.09it/s]

QAT: 100%|██████████| 5/5 [00:04<00:00,  1.03it/s]

int8 QAT final: (2.0550265312194824, 0.60107421875, 2.2567036151885986)

In [13]:

# Side-by-side: fp32 final vs int8 QAT final.
fp_final = history[-1]
qat_final = qat_history[-1]
print(f"fp32     val_acc={fp_final[1]:.4f}  val_loss={fp_final[2]:.4f}")
print(f"int8 QAT val_acc={qat_final[1]:.4f}  val_loss={qat_final[2]:.4f}")

# Weights have shrunk ~4x in storage; latency on accelerators with int8 fast paths
# (H100, TPU v4+) should also drop accordingly.

# Side-by-side: fp32 final vs int8 QAT final. fp_final = history[-1] qat_final = qat_history[-1] print(f"fp32 val_acc={fp_final[1]:.4f} val_loss={fp_final[2]:.4f}") print(f"int8 QAT val_acc={qat_final[1]:.4f} val_loss={qat_final[2]:.4f}") # Weights have shrunk ~4x in storage; latency on accelerators with int8 fast paths # (H100, TPU v4+) should also drop accordingly.

fp32     val_acc=0.5596  val_loss=2.2989
int8 QAT val_acc=0.6011  val_loss=2.2567

In [14]:

# BitNet b1.58-style: ternary weights (int2 fallback since qwix has no native
# 1.58-bit format yet) with int8 activations. Bigger accuracy dip than int8
# W+A; usually recovers with a few more fine-tune epochs.
bitnet_rules = spyx.quant.bitnet_ternary_rules()
bitmodel = spyx.quant.quantize(model, sample_events, rules=bitnet_rules)
bit_optimizer = nnx.Optimizer(bitmodel, optax.lion(1e-4), wrt=nnx.Param)

# One quick QAT epoch — extend for serious comparisons.
key, shuffle_key = jax.random.split(key)
perm = jax.random.permutation(shuffle_key, N_BATCHES)
for idx in perm:
    train_step(bitmodel, bit_optimizer, train_obs[idx], train_labels[idx])
bit_acc, bit_loss = eval_epoch(bitmodel, test_obs, test_labels)
print(f"BitNet ternary val_acc after 1 QAT epoch: {float(bit_acc):.4f}  val_loss: {float(bit_loss):.4f}")

# BitNet b1.58-style: ternary weights (int2 fallback since qwix has no native # 1.58-bit format yet) with int8 activations. Bigger accuracy dip than int8 # W+A; usually recovers with a few more fine-tune epochs. bitnet_rules = spyx.quant.bitnet_ternary_rules() bitmodel = spyx.quant.quantize(model, sample_events, rules=bitnet_rules) bit_optimizer = nnx.Optimizer(bitmodel, optax.lion(1e-4), wrt=nnx.Param) # One quick QAT epoch — extend for serious comparisons. key, shuffle_key = jax.random.split(key) perm = jax.random.permutation(shuffle_key, N_BATCHES) for idx in perm: train_step(bitmodel, bit_optimizer, train_obs[idx], train_labels[idx]) bit_acc, bit_loss = eval_epoch(bitmodel, test_obs, test_labels) print(f"BitNet ternary val_acc after 1 QAT epoch: {float(bit_acc):.4f} val_loss: {float(bit_loss):.4f}")

BitNet ternary val_acc after 1 QAT epoch: 0.0898  val_loss: 5.2334

Notes on this migration¶

• The original notebook used haiku.mixed_precision.set_policy(...) for fp16 weight casts. Flax NNX does not yet expose an equivalent first-class API; it has been omitted here. For best throughput on modern accelerators set jax.config.update("jax_default_matmul_precision", "bfloat16") at the top of the notebook.
• The old spyx.loaders module is now spyx.data; the constructor signature is unchanged but train_epoch() is a Python generator instead of a stacked tensor, so the outer loop uses Python iteration rather than jax.lax.scan.