Welcome to Spyx!

Spyx (pronounced "spikes") is a compact spiking neural network (SNN) library built on JAX and Flax NNX. It delivers the flexibility and extensibility of PyTorch-based SNN libraries while training at speeds comparable to — or faster than — frameworks that hand-write custom CUDA kernels, thanks to JAX's JIT compilation and auto-differentiation. Trained models can be exported to neuromorphic hardware via the Neuromorphic Intermediate Representation (NIR).

Be sure to give it a star on GitHub: kmheckel/spyx

The module map

Spyx is a stable core plus a clearly-fenced experimental staging area. Build on the core; reach into spyx.experimental deliberately, knowing its API can move without a deprecation cycle.

Stable core — public API, kept backwards-compatible:

Module	One-liner
spyx.nn	Spiking neurons (LIF, ALIF, CuBaLIF, LI, IF + recurrent variants), Sequential, Flatten, and the time-major run helper.
spyx.axn	Surrogate-gradient factories (SuperSpike, Arctan, Tanh, Boxcar, Triangular, straight-through).
spyx.fn	Losses & metrics — integral_crossentropy, integral_accuracy, silence_reg, sparsity_reg, mse_spikerate.
spyx.ssm	State-space layers — LRU, S5Diag, Mamba, ChunkedSSM.
spyx.phasor	Phasor networks — PhasorLinear/Activation/Readout/MLP, SpikingPhasor.
spyx.data	Grain event-data loaders with rate / angle / latency encoding (SHD_loader, NMNIST_loader).
spyx.optimize	Quick-training helpers — fit, make_train_step, make_eval_step.
spyx.nir	NIR import/export for neuromorphic-hardware interop.
spyx.quant	qwix-backed int8 / int4 / ternary quantization-aware training.
spyx.bench	Measurement — latency, throughput, peak memory, FLOP/MFU, spike-rate energy proxy.

Experimental — unstable API under spyx.experimental, the staging area for in-progress research:

Module	One-liner
spyx.experimental.PSU_LIF	Reset-free parallel LIF (associative-scan).
spyx.experimental.ResonateFire	Complex resonate-and-fire oscillatory neuron.
spyx.experimental.raven	Routing-slot memory (RavenRSM) + spiking sibling for high-recall sequence modeling.
spyx.experimental.compress	Bit-packed activation storage for memory-efficient BPTT.
spyx.experimental.stochastic	Stochastic (Bernoulli-spiking) & parallelizable prototypes (SPSN, ...).

Import experimental pieces from spyx.experimental so the dependency is explicit. See Research with Spyx for how work graduates from here into the core.

How the documentation is organised

These docs follow the Diátaxis framework: four sections, each answering a different kind of question.

Tutorials — learning-oriented

New to Spyx? Start here. Guided lessons that take you from an empty environment to a working, trained SNN, one step at a time.

• Your first SNN — install Spyx, build a spiking network, and train it on the Spiking Heidelberg Digits dataset.
• Notebook tutorials on surrogate gradients, neuroevolution, NIR conversion, quantization, state-space models, and phasor networks.

How-to guides — task-oriented

Need to get something done? Short, focused recipes for specific tasks, assuming you already know the basics.

• Install Spyx · Train a model · Quantize a model
• Export / import via NIR · Use SSM and phasor layers · Load event data

Reference — information-oriented

Looking up a signature? Exhaustive, auto-generated API documentation for every public module: spyx.nn, spyx.axn, spyx.fn, spyx.data, spyx.optimize, spyx.nir, spyx.quant, spyx.ssm, spyx.phasor, spyx.bench, spyx.experimental.

Explanation — understanding-oriented

Want to understand why? Background and design discussion, read away from the keyboard.

• Design and architecture — why Spyx is built the way it is, the module map, and how it compares to PyTorch SNN libraries.
• A primer on spiking neural networks — spikes, LIF dynamics, surrogate gradients, and rate vs. latency coding.
• Parallel spiking neurons — reset-free, associative-scan neurons that train in O(log T) depth.
• Research with Spyx — how the research/ corpus is organised and how to add a study.

For developers

If you're contributing to Spyx or working with AI coding agents, check out AGENTS.md for a comprehensive overview of the project structure, development workflow, and coding standards.