Every number is measured on an NVIDIA T4 and reproducible from benchmarks/, including the runs that went against us.
Speed or flexibility is a false choice
Libraries built on hand-written CUDA kernels are fast but only support the neuron models somebody already wrote a kernel for. Libraries in pure PyTorch or JAX let you write any neuron and run considerably slower.
jaxpike attacks that from the algorithmic side first. Parallel-in-time execution, rematerialization and online learning are all pure JAX, need no kernel code, and work on any neuron that fits a three-method contract.
Networks are ordinary pytrees, state is explicit and functional, and everything composes with jax.jit, jax.grad and jax.vmap.
import jax
import jaxpike as jp
k1, k2 = jax.random.split(jax.random.key(0))
net = jp.Sequential(
jp.Dense(700, 256, key=k1, gain=jp.lif_gain(20.0)),
jp.LinearLIF(tau=20.0, threshold=0.5),
jp.Dense(256, 20, key=k2, gain=jp.lif_gain(20.0)),
jp.LeakyIntegrator(tau=20.0),
)
membrane, state = jp.unroll_parallel(net, xs) # (time, batch, features)
logits = jp.max_membrane_logits(membrane)
My network will not train
Deep SNNs go silent: activity decays with depth until nothing fires and there is no gradient anywhere. One line fixes it.
Read the diagnosis →I am porting from snnTorch
Two conventions differ numerically — input normalization and reset timing — and both change results silently rather than raising.
Migration guide →I need it faster
Three execution paths with identical signatures, and an honest account of which speedup applies to what.
Execution guide →