# SNNAX -- Spiking Neural Networks in JAX

arXiv: https://arxiv.org/abs/2409.02842 · HTML: https://mlsanity.com/paper/2409.02842/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 16 of its 31 references that are in the map.
- First posted on arXiv 4 Sep 2024.

## Builds on (most cited first, 16 of 16)

- [Automatic differentiation in machine learning: a survey](https://mlsanity.com/paper/1502.05767/) (2015)
- [Surrogate Gradient Learning in Spiking Neural Networks](https://mlsanity.com/paper/1901.09948/) (2019)
- [On Neural Differential Equations](https://mlsanity.com/paper/2202.02435/) (2022)
- [SLAYER: Spike Layer Error Reassignment in Time](https://mlsanity.com/paper/1810.08646/) (2018)
- [Training Spiking Neural Networks Using Lessons From Deep Learning](https://mlsanity.com/paper/2109.12894/) (2021)
- [SpikingJelly: An open-source machine learning infrastructure platform for spike-based intelligence](https://mlsanity.com/paper/2310.16620/) (2023)
- [Equinox: neural networks in JAX via callable PyTrees and filtered transformations](https://mlsanity.com/paper/2111.00254/) (2021)
- [Event-Based Backpropagation can compute Exact Gradients for Spiking Neural Networks](https://mlsanity.com/paper/2009.08378/) (2020)
- [Brain-Inspired Learning on Neuromorphic Substrates](https://mlsanity.com/paper/2010.11931/) (2020)
- [NeuroBench: A Framework for Benchmarking Neuromorphic Computing Algorithms and Systems](https://mlsanity.com/paper/2304.04640/) (2023)
- [Neuromorphic Intermediate Representation: A Unified Instruction Set for Interoperable Brain-Inspired Computing](https://mlsanity.com/paper/2311.14641/) (2023)
- [Meta-learning Spiking Neural Networks with Surrogate Gradient Descent](https://mlsanity.com/paper/2201.10777/) (2022)
- [jaxsnn: Event-driven Gradient Estimation for Analog Neuromorphic Hardware](https://mlsanity.com/paper/2401.16841/) (2024)
- [EXODUS: Stable and Efficient Training of Spiking Neural Networks](https://mlsanity.com/paper/2205.10242/) (2022)
- [Slax: A Composable JAX Library for Rapid and Flexible Prototyping of Spiking Neural Networks](https://mlsanity.com/paper/2404.05807/) (2024)
- [Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks](https://mlsanity.com/paper/2402.18994/) (2024)

## Built on by (most cited first, 3 of 3)

- [Edge Intelligence with Spiking Neural Networks](https://mlsanity.com/paper/2507.14069/) (2025)
- [Training event-based neural networks with exact gradients via Differentiable ODE Solving in JAX](https://mlsanity.com/paper/2603.08146/) (2026)
- [A Truly Sparse and General Implementation of Gradient-Based Synaptic Plasticity](https://mlsanity.com/paper/2501.11407/) (2025)

## Abstract (arXiv)

Spiking Neural Networks (SNNs) simulators are essential tools to prototype biologically inspired models and neuromorphic hardware architectures and predict their performance. For such a tool, ease of use and flexibility are critical, but so is simulation speed especially given the complexity inherent to simulating SNN. Here, we present SNNAX, a JAX-based framework for simulating and training such models with PyTorch-like intuitiveness and JAX-like execution speed. SNNAX models are easily extended and customized to fit the desired model specifications and target neuromorphic hardware. Additionally, SNNAX offers key features for optimizing the training and deployment of SNNs such as flexible automatic differentiation and just-in-time compilation. We evaluate and compare SNNAX to other commonly used machine learning (ML) frameworks used for programming SNNs. We provide key performance metrics, best practices, documented examples for simulating SNNs in SNNAX, and implement several benchmarks used in the literature.