spyx
spyx copied to clipboard
kmheckel
Spyx: Spiking Neural Networks in JAX
Hi! I was working on https://www.nature.com/articles/s41467-020-17236-y, so I added two classes based on the original code of the paper: `RecurrentLIFLight` and `LeakyLinear` to use the same LIF/ALIF neurons and the...
There might be a bug with respect to time_major = False or True that could present incorrect results since scanning over one axis is faster than scanning over the other...
https://jax.readthedocs.io/en/latest/pallas/design.html Create pallas kernel for LIF neurons and investigate to see if it's faster than just static unrolling the LIF as defined in Spyx.
It would be cool to see how synthetic gradients could fit into the SNN training schema: https://arxiv.org/abs/1608.05343 https://greydanus.github.io/2016/11/26/synthetic-gradients/
Make spyx.optimize: - [ ] surrogate gradient training loop - [ ] take SNN, params, Dataset in form (#batches, batchsize, timesteps, ...), optimizer - [ ] also take number of...
Create functions to load/save models to HDF5 under the Neuromorphic Intermediate Representation standard to facilitate cross-platform deployment. https://nnir.readthedocs.io/en/latest/what.html - [x] Implement spyx.nir.to_nir() - [x] Implement spyx.nir.from_nir() - [x] Support feed-forward...
The current documentation is passable but it could always be better. - [x] Make spyx.nn listings cleaner by removing unwanted special-members - [x] Incorporate notebooks from the examples/research folders. -...
It's a common mistake to get the output shape wrong compared to the target label shape, adding some kind of check in the loss and acc functions to throw an...
Right now there's no test cases to verify that changes to the code base work for actually training models/there's no way to detect if changes to other packages might break...
Use Chex to add type checking to the library in spots not already handled by Haiku.