TY  - CONF
AU  - Lohoff, Jamie
AU  - Finkbeiner, Jan
AU  - Neftci, Emre
TI  - SNNAX-Spiking Neural Networks in JAX
M1  - FZJ-2025-01092
SP  - 251 - 255
PY  - 2024
AB  - 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.
T2  - 2024 International Conference on Neuromorphic Systems (ICONS)
CY  - 30 Jul 2024 - 2 Aug 2024, Arlington, Virginia (USA)
Y2  - 30 Jul 2024 - 2 Aug 2024
M2  - Arlington, Virginia, USA
LB  - PUB:(DE-HGF)8
DO  - DOI:10.34734/FZJ-2025-01092
UR  - https://juser.fz-juelich.de/record/1038043
ER  -