Spiking Neural Networks implemented on top of diffrax. Features include:
- simulating trajectories of Leaky-Integrate-and-Fire neurons;
- stochastic firing through intensity functions;
- Stochastic Spiking Neural Networks (SSNNs) as introduced here;
- arbitrary network structures;
- automatic differentiation of spike times and neuronal state variables.
This project is still in a very early experimental phase. For now everything should work more or less out of the box, but I am not sure how much I will maintain this project in the future.
One thing to note is that the latest relase of jax (v0.6.0) introduced a bug that manifests in the current implementation. For now, installing this repository will also install v0.5.3, but I will probably update the dependencies in the future, when the fix is included.
Clone the repository:
git clone https://github.com/cholberg/snnax
cd snnax
and install:
pip install .
For some usage examples see the example notebook. To reproduce the results of Exact Gradients for Stochastic Spiking Neural Networks Driven by Rough Signals simply run the notebooks here and here.