⚡🧠💻 Welcome to Spyx! 💻🧠⚡

Why use Spyx?

Spyx (pronounced "spikes") is a compact spiking neural network library built on JAX and Flax NNX, offering the flexibility and extensibility of PyTorch-based frameworks while enabling the extreme performance of SNN libraries which implement custom CUDA kernels for their dynamics.

The library currently supports training SNNs via surrogate gradient descent and neuroevolution, with additional capabilities such as ANN2SNN conversion and Phasor Networks being planned for the future. Spyx offers a number of predefined neuron models but is designed for it to be easy to define your own and plug it into a model; the hope is to soon include definitions of SpikingRWKV and other more sophisticated model blocks into the framework.

Installation:

Spyx now uses uv for package management.

To install Spyx:

uv add spyx

If you need the data loading utilities:

uv add spyx[loaders]

Note: As with other libraries built on top of JAX, you need to install jax with GPU if you want to get the full benefit of this library. Directions for installing JAX with GPU support can be found at the following: https://github.com/google/jax#installation

Hardware Requirements:

Spyx achieves extremely high performance by maintaining the entire dataset in the GPU's vRAM; as such a decent amount of memory for both the CPU and GPU are needed to handle the dataset loading and then training. For smaller networks of only several hundred thousand parameters, the training process can be comfortably executed on even laptop GPU's with only 6GB of vRAM. For large SNNs or for neuroevolution it is recommended to use a higher memory card.

Since Spyx is developed on top of the current JAX version, it does not work on Google Colab's TPUs which use an older version. Cloud TPU support will be tested in the near future.

Development:

To set up a development environment:

git clone https://github.com/kmheckel/spyx
cd spyx
uv sync

This will install all dependencies including development tools (pytest, ruff, mkdocs).

End-to-end install check

Once installed, run the bundled install-check script to confirm JAX, Spyx, and optional extras are all wired up correctly:

uv run python scripts/check_install.py

Seven checks in ~30 seconds — JAX version + visible devices, Spyx imports, SNN forward pass, one training epoch, NIR roundtrip, notebook-API smoke tests, and optional-extra detection (tonic, qwix). Useful right after uv sync on a new machine, especially if you expect GPU / TPU devices to show up.

Code Quality

Spyx uses Ruff for linting and code formatting. Before committing changes, run:

# Check for linting errors
uv run ruff check

# Auto-fix linting errors
uv run ruff check --fix

# Format code
uv run ruff format

Ruff is configured in pyproject.toml to enforce code quality standards including:

• Import sorting (isort-compatible)
• PEP 8 style guidelines
• Common bug patterns (flake8-bugbear)
• Exclusion of research, docs, and scripts directories

Testing

Run the test suite using pytest:

# Run all tests
uv run pytest

# Run tests with verbose output
uv run pytest -v

# Run a specific test file
uv run pytest tests/test_data_grain.py

# Run tests with coverage
uv run pytest --cov=spyx --cov-report=html

Tests are located in the tests/ directory and cover core functionality including data loading, neuron models, and training utilities.

Releasing new versions

A utility script is provided to automate the release process:

# Dry run to see what would happen
python scripts/release.py --dry-run --github --pypi

# Perform a full release
python scripts/release.py --github --pypi

This script will:

1. Build the package using uv build.
2. Create a git tag and GitHub release using gh.
3. Publish to PyPI using uv publish.

Building Documentation

The documentation is built using MkDocs:

# Preview documentation locally
uv run mkdocs serve

# Build documentation
uv run mkdocs build

Research and Projects Using Spyx:

Experiments/Benchmarks used in the Spyx Paper: Benchmark Notebooks

Master's Thesis: Neuroevolution of Spiking Neural Networks

*** Your projects and research could be here! ***

Note: notebooks under research/ predate the Flax NNX migration and still use the legacy Haiku API. They are kept for reproducibility of the Spyx paper experiments and will be ported in a follow-up effort. New work should use the tutorials under docs/examples/.

Contributing:

If you'd like to contribute, head on over to the issues page to find proposed enhancements and leave a comment! Also head over to the Open Neuromorphic Discord server to ask questions!

Citation:

If you find Spyx useful in your work please cite it using the following Bibtex entries:

@misc{heckel2024spyx,
    title={Spyx: A Library for Just-In-Time Compiled Optimization of Spiking Neural Networks},
    author={Kade M. Heckel and Thomas Nowotny},
    year={2024},
    eprint={2402.18994},
    archivePrefix={arXiv},
    primaryClass={cs.NE}
}

@software{kade_heckel_2024_10635178,
  author       = {Kade Heckel and
                  Steven Abreu and
                  Gregor Lenz and
                  Thomas Nowotny},
  title        = {kmheckel/spyx: v0.1.17},
  month        = feb,
  year         = 2024,
  publisher    = {Zenodo},
  version      = {camera-ready},
  doi          = {10.5281/zenodo.10635178},
  url          = {https://doi.org/10.5281/zenodo.10635178}
}