Eventax provides a JAX implementation of the EventProp algorithm using Diffrax ODE-solvers and Equinox offering full autograd support and easy extension with custom neuron dynamics.

Features

• Fully differentiable implementation via JAX and Diffrax
• Easy extension with custom neuron model dynamics + learnable parameters
• Support for (trainable) synnaptic delays.
• GPU/TPU compatibility through JAX

📦 Installation

pip install eventax

Cite as:

@misc{könig2026trainingeventbasedneuralnetworks,
      title={Training event-based neural networks with exact gradients via Differentiable ODE Solving in JAX}, 
      author={Lukas König and Manuel Kuhn and David Kappel and Anand Subramoney},
      year={2026},
      eprint={2603.08146},
      archivePrefix={arXiv},
      primaryClass={cs.LG},
      url={https://arxiv.org/abs/2603.08146}, 
}

You can read the paper here

Documentation:

Look at the documentation here.