001038043 001__ 1038043 001038043 005__ 20250203103306.0 001038043 0247_ $$2datacite_doi$$a10.34734/FZJ-2025-01092 001038043 037__ $$aFZJ-2025-01092 001038043 1001_ $$0P:(DE-Juel1)192147$$aLohoff, Jamie$$b0$$ufzj 001038043 1112_ $$a2024 International Conference on Neuromorphic Systems (ICONS)$$cArlington, Virginia$$d2024-07-30 - 2024-08-02$$wUSA 001038043 245__ $$aSNNAX-Spiking Neural Networks in JAX 001038043 260__ $$c2024 001038043 300__ $$a251 - 255 001038043 3367_ $$2ORCID$$aCONFERENCE_PAPER 001038043 3367_ $$033$$2EndNote$$aConference Paper 001038043 3367_ $$2BibTeX$$aINPROCEEDINGS 001038043 3367_ $$2DRIVER$$aconferenceObject 001038043 3367_ $$2DataCite$$aOutput Types/Conference Paper 001038043 3367_ $$0PUB:(DE-HGF)8$$2PUB:(DE-HGF)$$aContribution to a conference proceedings$$bcontrib$$mcontrib$$s1738246115_31341 001038043 520__ $$aSpiking 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. 001038043 536__ $$0G:(DE-HGF)POF4-5234$$a5234 - Emerging NC Architectures (POF4-523)$$cPOF4-523$$fPOF IV$$x0 001038043 536__ $$0G:(EU-Grant)953775$$aGREENEDGE - Taming the environmental impact of mobile networks through GREEN EDGE computing platforms (953775)$$c953775$$fH2020-MSCA-ITN-2020$$x1 001038043 536__ $$0G:(BMBF)03ZU1106CA$$aBMBF 03ZU1106CA - NeuroSys: Algorithm-Hardware Co-Design (Projekt C) - A (03ZU1106CA)$$c03ZU1106CA$$x2 001038043 536__ $$0G:(DE-Juel1)BMBF-03ZU1106CB$$aBMBF 03ZU1106CB - NeuroSys: Algorithm-Hardware Co-Design (Projekt C) - B (BMBF-03ZU1106CB)$$cBMBF-03ZU1106CB$$x3 001038043 536__ $$0G:(BMBF)16ME0400$$aBMBF 16ME0400 - Verbundprojekt: Neuro-inspirierte Technologien der künstlichen Intelligenz für die Elektronik der Zukunft - NEUROTEC II - (16ME0400)$$c16ME0400$$x4 001038043 7001_ $$0P:(DE-Juel1)190112$$aFinkbeiner, Jan$$b1$$ufzj 001038043 7001_ $$0P:(DE-Juel1)188273$$aNeftci, Emre$$b2$$ufzj 001038043 8564_ $$uhttps://ieeexplore.ieee.org/document/10766537 001038043 8564_ $$uhttps://juser.fz-juelich.de/record/1038043/files/SNNAX-Spiking%20Neural%20Networks%20in%20JAX.pdf$$yOpenAccess 001038043 909CO $$ooai:juser.fz-juelich.de:1038043$$pdnbdelivery$$pec_fundedresources$$pVDB$$pdriver$$popen_access$$popenaire 001038043 9101_ $$0I:(DE-588b)5008462-8$$6P:(DE-Juel1)192147$$aForschungszentrum Jülich$$b0$$kFZJ 001038043 9101_ $$0I:(DE-588b)5008462-8$$6P:(DE-Juel1)190112$$aForschungszentrum Jülich$$b1$$kFZJ 001038043 9101_ $$0I:(DE-588b)5008462-8$$6P:(DE-Juel1)188273$$aForschungszentrum Jülich$$b2$$kFZJ 001038043 9131_ $$0G:(DE-HGF)POF4-523$$1G:(DE-HGF)POF4-520$$2G:(DE-HGF)POF4-500$$3G:(DE-HGF)POF4$$4G:(DE-HGF)POF$$9G:(DE-HGF)POF4-5234$$aDE-HGF$$bKey Technologies$$lNatural, Artificial and Cognitive Information Processing$$vNeuromorphic Computing and Network Dynamics$$x0 001038043 9141_ $$y2024 001038043 915__ $$0StatID:(DE-HGF)0510$$2StatID$$aOpenAccess 001038043 920__ $$lyes 001038043 9201_ $$0I:(DE-Juel1)PGI-15-20210701$$kPGI-15$$lNeuromorphic Software Eco System$$x0 001038043 980__ $$acontrib 001038043 980__ $$aVDB 001038043 980__ $$aUNRESTRICTED 001038043 980__ $$aI:(DE-Juel1)PGI-15-20210701 001038043 9801_ $$aFullTexts