sparxml 0.1.0__tar.gz

sparxml-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024-2026 Ashish Kumar Singh
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

sparxml-0.1.0/PKG-INFO

@@ -0,0 +1,230 @@
+Metadata-Version: 2.4
+Name: sparxml
+Version: 0.1.0
+Summary: Spiking neural networks in JAX and Flax
+Author-email: Ashish Kumar Singh <ashishkmr472@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://sparxml.dev
+Project-URL: Repository, https://github.com/AshishKumar4/sparx
+Project-URL: Documentation, https://sparxml.dev/docs/
+Project-URL: Issues, https://github.com/AshishKumar4/sparx/issues
+Project-URL: Changelog, https://github.com/AshishKumar4/sparx/blob/main/CHANGELOG.md
+Keywords: jax,flax,spiking-neural-networks,neuromorphic,computational-neuroscience,connectomics,dew
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Intended Audience :: Developers
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Typing :: Typed
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: dewml<0.2,>=0.1.0
+Requires-Dist: jax
+Requires-Dist: flax
+Requires-Dist: optax
+Requires-Dist: numpy
+Provides-Extra: cuda12
+Requires-Dist: dewml[cuda12]<0.2,>=0.1.0; extra == "cuda12"
+Provides-Extra: cuda13
+Requires-Dist: dewml[cuda13]<0.2,>=0.1.0; extra == "cuda13"
+Provides-Extra: tpu
+Requires-Dist: dewml[tpu]<0.2,>=0.1.0; extra == "tpu"
+Provides-Extra: datasets
+Requires-Dist: h5py; extra == "datasets"
+Provides-Extra: connectome
+Requires-Dist: pyarrow; extra == "connectome"
+Provides-Extra: nir
+Requires-Dist: nir; extra == "nir"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: ruff==0.14.3; extra == "test"
+Requires-Dist: h5py; extra == "test"
+Requires-Dist: nir; extra == "test"
+Requires-Dist: pyarrow; extra == "test"
+Dynamic: license-file
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/banner-dark.svg">
+  <img alt="sparx: spiking neural networks in JAX" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/banner-light.svg" width="100%">
+</picture>
+
+sparx trains spiking neural networks and simulates circuits of biological neurons, in JAX. Its spiking layers are Flax modules, so they train with optax or [dew](https://github.com/AshishKumar4/dew) and work with `jit`, `grad`, `vmap` and sharding. The same neuron models also run in millivolts and milliseconds, wired into circuits and whole connectomes, and there they match NEST and Brian2.
+
+[sparxml.dev](https://sparxml.dev): a course from one neuron to a spiking network that drives from events, the docs and the API reference.
+
+[Guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md) · [Train, serve and export](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/train-and-deploy.md) · [From NEST and Brian2](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/nest-and-brian2.md) · [Fit a circuit](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/fit-a-circuit.md) · [Units](https://github.com/AshishKumar4/sparx/blob/main/docs/units.md) · [Status](https://github.com/AshishKumar4/sparx/blob/main/docs/status.md) · [Fidelity ledger](https://github.com/AshishKumar4/sparx/blob/main/docs/fidelity.md) · [Design](https://github.com/AshishKumar4/sparx/blob/main/docs/design.md) · [Performance](https://github.com/AshishKumar4/sparx/blob/main/docs/performance.md)
+
+## Install
+
+sparx is published as `sparxml` and imported as `sparx`, as [dew](https://github.com/AshishKumar4/dew) is published as `dewml` and imported as `dew`:
+
+| Where | Install |
+| --- | --- |
+| CPU | `pip install sparxml` |
+| NVIDIA GPU (CUDA 12 or 13) | `pip install "sparxml[cuda12]"` or `"sparxml[cuda13]"` |
+| TPU | `pip install "sparxml[tpu]"` |
+| SHD and MNIST readers, connectome tables, NIR | add the extras `datasets`, `connectome` and `nir` |
+
+From a clone, on the dew commit and the jax build the tests run on:
+
+```bash
+git clone https://github.com/AshishKumar4/sparx.git && cd sparx
+uv venv --python 3.12 && source .venv/bin/activate
+uv pip install -e ".[datasets,test]" -c constraints.txt && pytest -q
+```
+
+sparx needs Python 3.12 or later, and CI tests it on 3.12, 3.13 and 3.14 on Linux and on 3.12 on macOS, with JAX 0.11.2 and Flax 0.12 on CPU. The API may change before 1.0.
+
+## A first network
+
+```python
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import optax
+
+import sparx
+
+
+class Net(nn.Module):
+    @nn.compact
+    def __call__(self, spikes):                     # [T, B, 784]
+        x = sparx.nn.LIF(tau=2.0)(nn.Dense(256)(spikes))
+        return sparx.nn.LI(tau=2.0)(nn.Dense(10)(x))  # membrane [T, B, 10]
+
+
+net = Net()
+images = jax.random.uniform(jax.random.key(0), (32, 784))  # intensities in [0, 1]
+labels = jnp.zeros(32, jnp.int32)
+spikes = sparx.encode.RateEncoder(steps=8)(jax.random.key(1), images)  # [8, 32, 784]
+params = net.init(jax.random.key(2), spikes)
+
+
+def loss(params):
+    logits = jnp.mean(net.apply(params, spikes), axis=0)  # mean membrane over time
+    return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
+
+
+grads = jax.grad(loss)(params)
+```
+
+`LIF` turns input currents into spikes, exactly 0 or 1, and `LI` integrates them into a membrane, which is the readout. Everything else is Flax and optax. [`examples/train_mnist.py`](https://github.com/AshishKumar4/sparx/blob/main/examples/train_mnist.py) trains a network like this on MNIST under dew's `Trainer`.
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-dark.webp">
+  <img alt="A spiking classifier learning MNIST: a test digit, the hidden layer's spikes for it, the output spike counts and the test accuracy rising over 400 training steps" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-light.webp" width="100%">
+</picture>
+
+A 784-200-10 network of LIF neurons learning MNIST by surrogate gradients. It follows one test digit through 400 training steps, showing the hidden layer's spikes, the ten output neurons' spike counts and the test accuracy, which reaches 92.8%.
+
+## How sparx fits together
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/architecture-dark.svg">
+  <img alt="Training tools and simulation tools both run neuron models through one protocol" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/architecture-light.svg" width="100%">
+</picture>
+
+Every neuron model implements one protocol, `init_state` and `step`, and `run` scans a model over time. Dimensionless cells serve deep learning, and physical models in mV and ms serve neuroscience. The two halves mix: a layer can hold a physical model (`nn.Dynamics(AdEx())`), and a simulated population can hold a dimensionless cell.
+
+## A spiking layer over time
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/over_time-dark.svg">
+  <img alt="A layer scans one step function over time; a LIF membrane rises to threshold and resets at each spike; the spike's gradient is a smooth surrogate" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/over_time-light.svg" width="100%">
+</picture>
+
+Arrays are time-major, `[T, B, ...]`. Synaptic layers run over all steps in one matrix product, and only the neurons step through time. A spike is a step function with zero derivative almost everywhere, so the backward pass uses a surrogate's slope instead. The layers are LIF, IF, LI, current-based synaptic LIF, adaptive LIF (ALIF), rate units, parallel spiking neurons and dense layers with learned delays ([guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#neuron-layers)).
+
+## Recurrence and fast weights
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/recurrence-dark.svg">
+  <img alt="A recurrent cell sends its output back through a dense, sparse or delayed wiring, with optional fast weights from a Hebbian trace" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/recurrence-light.svg" width="100%">
+</picture>
+
+`RecurrentCell` feeds any model's output back through a wiring: dense, an edge list such as a connectome's, or edges with their own delays. Fast weights add a Hebbian trace that each sequence writes as it runs (Miconi et al. 2018, 2019). Against Miconi et al.'s four networks in PyTorch, activity, traces and gradients agree within 5e-14. On their pattern completion task, the plastic network gets 0.3% of the zeroed bits wrong, and the same network without fast weights 50.1%.
+
+## Learning rules
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/learning-dark.svg">
+  <img alt="Nine ways to train a spiking network in sparx, each with a schematic of its learning signal" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/learning-light.svg" width="100%">
+</picture>
+
+Each rule is checked against what defines it. e-prop meets the two identities its authors verify their code with, OTTT matches their PyTorch modules, PC-ALM matches their JAX reference to 5e-14, and conversion matches their toolbox. REINFORCE is checked on enumerated trajectories, and exact spike times against finite differences. The [guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#learning-rules) describes each rule.
+
+## Simulating circuits
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/circuits-dark.svg">
+  <img alt="Two populations connected by excitatory and inhibitory projections with delays, driven by Poisson input and simulated in chunks" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/circuits-light.svg" width="100%">
+</picture>
+
+```python
+import jax
+from sparx.graph import PopulationRate, SpikeRaster, simulate
+from sparx.graph.models import brunel
+
+network = brunel(250, g=5.0, eta=2.0)        # 1,250 LIF neurons; brunel(2500) is the paper's 12,500
+result = simulate(network, network.init(jax.random.key(0)), duration=200.0, key=jax.random.key(1),
+                  monitors={"spikes": SpikeRaster("e"), "rate": PopulationRate("e")})
+spikes = result.records["spikes"]             # [2000, 1000]: one row of booleans per 0.1 ms step
+```
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/network-dark.webp">
+  <img alt="A raster of 200 neurons of Brunel's network firing irregularly over 300 ms, with the population rate below" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/network-light.webp" width="100%">
+</picture>
+
+Brunel's balanced network at the paper's size, 10,000 excitatory and 2,500 inhibitory LIF neurons, in its asynchronous irregular regime at 37 Hz.
+
+The physical models match NEST 3.10 and Brian2 2.10 spike for spike where the dynamics are deterministic, for the integration scheme, dtype and step each check states ([status](https://github.com/AshishKumar4/sparx/blob/main/docs/status.md#capabilities-and-limits)), and in rate, irregularity and synchrony where they are chaotic. Potjans and Diesmann's cortical microcircuit, built as its reference implementation builds it, fires spike for spike with NEST on the same network ([from NEST and Brian2](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/nest-and-brian2.md)). On a 4-core CPU, sparx simulates a second of Brunel's network in 9.6 s, NEST in 7.5 s and Brian2 in 11.8 s ([performance](https://github.com/AshishKumar4/sparx/blob/main/docs/performance.md#against-nest-and-brian2)). Populations can hold graded neurons and connect through stochastic release, gap junctions and neuromodulators. Projections can carry STDP, triplet STDP, dopamine-modulated STDP and short-term plasticity ([guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#simulating-circuits)).
+
+## Connectomes
+
+`sparx.graph.connectome` builds Shiu et al.'s (2024) model of the whole fly brain from FlyWire. It reproduces their published runs, with a rate correlation of 0.999 and the motor neuron MN9 at 67.1 Hz against their 67.0 ± 6.6, at about 30 s per simulated second on 4 CPU cores. `FLYNN` (Wang and Chen 2026) trains a connectome as a recurrent rate network with one learned weight per synapse; against their PyTorch cell its activity and gradients agree within 1e-15.
+
+## RNeuralNet
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/messages-dark.webp">
+  <img alt="Messages travelling along the connections of a small RNeuralNet, each connection with its own delay" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/messages-light.webp" width="100%">
+</picture>
+
+`sparx.learn.RNeuralNet` rebuilds RNeuralNet-Research (2018), an early project of the author's, deterministically. Graded neurons sit on a random graph, each connection delivers its messages after its own delay, and a reward spreads backward by a softmax of activity. Compiled and run in a fixed order, the original C++ and sparx agree within 7.2e-7. On a delayed cue-order task, REINFORCE through the same network learns the task on four of five seeds, and the reward-diffusion rule never changes the network's choice. AGREL's update, a signed error sent back from the chosen output through the weights, learns it on the same four seeds; the same error spread by the original's shares does not.
+
+## Training on dew
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-dark.svg">
+  <img alt="A Flax model and a sparx objective go to dew's trainer, which writes a run record that reloads, serves and exports" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-light.svg" width="100%">
+</picture>
+
+The `Trainer` from [dew](https://github.com/AshishKumar4/dew) runs sparx's objectives for classification, activity fitting, e-prop, predictive coding and RNeuralNet's rewards. A run's record names every class by import path, so `dew.pipeline("runs/shd", trust=("sparx",))` loads a trained network in a new process, and `sparx.serve.StreamServer` serves it to many streams at once. The [guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#training-on-dew) has a full SHD script.
+
+## Results
+
+| Task | Network | Test accuracy |
+| --- | --- | --- |
+| MNIST, rate-coded, 8 steps | 784-512-512 LIF | 97.5% after 2 epochs |
+| SHD, Hammouamri et al.'s recipe, 150 epochs, three seeds | 140-256-256 LIF with learned delays | 93.99 ± 0.29% at the last epoch (their code on the same GPU: 93.89 ± 0.26%) |
+| SHD, 140 channels | 140-128 ALIF, with and without learned delays | 74.6% and 64.5% |
+| Fashion-MNIST, Seely and Gould's headline cell | ReLU residual MLP, depth 32 | PC-ALM 75.1%, PC 62.2%, backpropagation 77.8% |
+| Pattern completion, Miconi et al.'s task | plastic recurrent network | 0.3% of bits wrong; 50.1% without fast weights |
+
+The SHD row is the full recipe beside the authors' code, both on an A100, three seeds each ([research/shd](https://github.com/AshishKumar4/sparx/blob/main/research/shd/README.md)). Both train on every training recording and score the test set after each epoch. The paper's 95.07 ± 0.24% (a 95% confidence interval over ten runs) is the best epoch on the test set, which chooses with the test set; here, as mean ± standard deviation over three seeds, their code's best epoch is 95.17 ± 0.61% and sparx's 94.96 ± 0.89%. With a tenth of the training set held out to choose the epoch, sparx scores 94.14 ± 0.98% on test. The other rows are short, untuned runs on a 4-core CPU; the [guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#results-in-detail) gives their commands, times and comparisons.
+
+## Correctness
+
+Every model is checked against a reference: a float64 loop of its equations, the original authors' code, or NEST and Brian2. [docs/fidelity.md](https://github.com/AshishKumar4/sparx/blob/main/docs/fidelity.md) lists each model's reference, the check, the observed error and every known difference. `pytest -q` runs all of it on CPU in about 16 minutes.
+
+[`tools/make_figures.py`](https://github.com/AshishKumar4/sparx/blob/main/tools/make_figures.py) draws the banner and diagrams, and [`tools/make_clips.py`](https://github.com/AshishKumar4/sparx/blob/main/tools/make_clips.py) renders the clips. The spikes in the banner and the clips come from sparx runs.
+
+## License
+
+MIT

sparxml-0.1.0/README.md

@@ -0,0 +1,180 @@
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/banner-dark.svg">
+  <img alt="sparx: spiking neural networks in JAX" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/banner-light.svg" width="100%">
+</picture>
+
+sparx trains spiking neural networks and simulates circuits of biological neurons, in JAX. Its spiking layers are Flax modules, so they train with optax or [dew](https://github.com/AshishKumar4/dew) and work with `jit`, `grad`, `vmap` and sharding. The same neuron models also run in millivolts and milliseconds, wired into circuits and whole connectomes, and there they match NEST and Brian2.
+
+[sparxml.dev](https://sparxml.dev): a course from one neuron to a spiking network that drives from events, the docs and the API reference.
+
+[Guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md) · [Train, serve and export](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/train-and-deploy.md) · [From NEST and Brian2](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/nest-and-brian2.md) · [Fit a circuit](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/fit-a-circuit.md) · [Units](https://github.com/AshishKumar4/sparx/blob/main/docs/units.md) · [Status](https://github.com/AshishKumar4/sparx/blob/main/docs/status.md) · [Fidelity ledger](https://github.com/AshishKumar4/sparx/blob/main/docs/fidelity.md) · [Design](https://github.com/AshishKumar4/sparx/blob/main/docs/design.md) · [Performance](https://github.com/AshishKumar4/sparx/blob/main/docs/performance.md)
+
+## Install
+
+sparx is published as `sparxml` and imported as `sparx`, as [dew](https://github.com/AshishKumar4/dew) is published as `dewml` and imported as `dew`:
+
+| Where | Install |
+| --- | --- |
+| CPU | `pip install sparxml` |
+| NVIDIA GPU (CUDA 12 or 13) | `pip install "sparxml[cuda12]"` or `"sparxml[cuda13]"` |
+| TPU | `pip install "sparxml[tpu]"` |
+| SHD and MNIST readers, connectome tables, NIR | add the extras `datasets`, `connectome` and `nir` |
+
+From a clone, on the dew commit and the jax build the tests run on:
+
+```bash
+git clone https://github.com/AshishKumar4/sparx.git && cd sparx
+uv venv --python 3.12 && source .venv/bin/activate
+uv pip install -e ".[datasets,test]" -c constraints.txt && pytest -q
+```
+
+sparx needs Python 3.12 or later, and CI tests it on 3.12, 3.13 and 3.14 on Linux and on 3.12 on macOS, with JAX 0.11.2 and Flax 0.12 on CPU. The API may change before 1.0.
+
+## A first network
+
+```python
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import optax
+
+import sparx
+
+
+class Net(nn.Module):
+    @nn.compact
+    def __call__(self, spikes):                     # [T, B, 784]
+        x = sparx.nn.LIF(tau=2.0)(nn.Dense(256)(spikes))
+        return sparx.nn.LI(tau=2.0)(nn.Dense(10)(x))  # membrane [T, B, 10]
+
+
+net = Net()
+images = jax.random.uniform(jax.random.key(0), (32, 784))  # intensities in [0, 1]
+labels = jnp.zeros(32, jnp.int32)
+spikes = sparx.encode.RateEncoder(steps=8)(jax.random.key(1), images)  # [8, 32, 784]
+params = net.init(jax.random.key(2), spikes)
+
+
+def loss(params):
+    logits = jnp.mean(net.apply(params, spikes), axis=0)  # mean membrane over time
+    return optax.softmax_cross_entropy_with_integer_labels(logits, labels).mean()
+
+
+grads = jax.grad(loss)(params)
+```
+
+`LIF` turns input currents into spikes, exactly 0 or 1, and `LI` integrates them into a membrane, which is the readout. Everything else is Flax and optax. [`examples/train_mnist.py`](https://github.com/AshishKumar4/sparx/blob/main/examples/train_mnist.py) trains a network like this on MNIST under dew's `Trainer`.
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-dark.webp">
+  <img alt="A spiking classifier learning MNIST: a test digit, the hidden layer's spikes for it, the output spike counts and the test accuracy rising over 400 training steps" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-light.webp" width="100%">
+</picture>
+
+A 784-200-10 network of LIF neurons learning MNIST by surrogate gradients. It follows one test digit through 400 training steps, showing the hidden layer's spikes, the ten output neurons' spike counts and the test accuracy, which reaches 92.8%.
+
+## How sparx fits together
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/architecture-dark.svg">
+  <img alt="Training tools and simulation tools both run neuron models through one protocol" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/architecture-light.svg" width="100%">
+</picture>
+
+Every neuron model implements one protocol, `init_state` and `step`, and `run` scans a model over time. Dimensionless cells serve deep learning, and physical models in mV and ms serve neuroscience. The two halves mix: a layer can hold a physical model (`nn.Dynamics(AdEx())`), and a simulated population can hold a dimensionless cell.
+
+## A spiking layer over time
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/over_time-dark.svg">
+  <img alt="A layer scans one step function over time; a LIF membrane rises to threshold and resets at each spike; the spike's gradient is a smooth surrogate" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/over_time-light.svg" width="100%">
+</picture>
+
+Arrays are time-major, `[T, B, ...]`. Synaptic layers run over all steps in one matrix product, and only the neurons step through time. A spike is a step function with zero derivative almost everywhere, so the backward pass uses a surrogate's slope instead. The layers are LIF, IF, LI, current-based synaptic LIF, adaptive LIF (ALIF), rate units, parallel spiking neurons and dense layers with learned delays ([guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#neuron-layers)).
+
+## Recurrence and fast weights
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/recurrence-dark.svg">
+  <img alt="A recurrent cell sends its output back through a dense, sparse or delayed wiring, with optional fast weights from a Hebbian trace" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/recurrence-light.svg" width="100%">
+</picture>
+
+`RecurrentCell` feeds any model's output back through a wiring: dense, an edge list such as a connectome's, or edges with their own delays. Fast weights add a Hebbian trace that each sequence writes as it runs (Miconi et al. 2018, 2019). Against Miconi et al.'s four networks in PyTorch, activity, traces and gradients agree within 5e-14. On their pattern completion task, the plastic network gets 0.3% of the zeroed bits wrong, and the same network without fast weights 50.1%.
+
+## Learning rules
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/learning-dark.svg">
+  <img alt="Nine ways to train a spiking network in sparx, each with a schematic of its learning signal" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/learning-light.svg" width="100%">
+</picture>
+
+Each rule is checked against what defines it. e-prop meets the two identities its authors verify their code with, OTTT matches their PyTorch modules, PC-ALM matches their JAX reference to 5e-14, and conversion matches their toolbox. REINFORCE is checked on enumerated trajectories, and exact spike times against finite differences. The [guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#learning-rules) describes each rule.
+
+## Simulating circuits
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/circuits-dark.svg">
+  <img alt="Two populations connected by excitatory and inhibitory projections with delays, driven by Poisson input and simulated in chunks" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/circuits-light.svg" width="100%">
+</picture>
+
+```python
+import jax
+from sparx.graph import PopulationRate, SpikeRaster, simulate
+from sparx.graph.models import brunel
+
+network = brunel(250, g=5.0, eta=2.0)        # 1,250 LIF neurons; brunel(2500) is the paper's 12,500
+result = simulate(network, network.init(jax.random.key(0)), duration=200.0, key=jax.random.key(1),
+                  monitors={"spikes": SpikeRaster("e"), "rate": PopulationRate("e")})
+spikes = result.records["spikes"]             # [2000, 1000]: one row of booleans per 0.1 ms step
+```
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/network-dark.webp">
+  <img alt="A raster of 200 neurons of Brunel's network firing irregularly over 300 ms, with the population rate below" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/network-light.webp" width="100%">
+</picture>
+
+Brunel's balanced network at the paper's size, 10,000 excitatory and 2,500 inhibitory LIF neurons, in its asynchronous irregular regime at 37 Hz.
+
+The physical models match NEST 3.10 and Brian2 2.10 spike for spike where the dynamics are deterministic, for the integration scheme, dtype and step each check states ([status](https://github.com/AshishKumar4/sparx/blob/main/docs/status.md#capabilities-and-limits)), and in rate, irregularity and synchrony where they are chaotic. Potjans and Diesmann's cortical microcircuit, built as its reference implementation builds it, fires spike for spike with NEST on the same network ([from NEST and Brian2](https://github.com/AshishKumar4/sparx/blob/main/docs/tutorials/nest-and-brian2.md)). On a 4-core CPU, sparx simulates a second of Brunel's network in 9.6 s, NEST in 7.5 s and Brian2 in 11.8 s ([performance](https://github.com/AshishKumar4/sparx/blob/main/docs/performance.md#against-nest-and-brian2)). Populations can hold graded neurons and connect through stochastic release, gap junctions and neuromodulators. Projections can carry STDP, triplet STDP, dopamine-modulated STDP and short-term plasticity ([guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#simulating-circuits)).
+
+## Connectomes
+
+`sparx.graph.connectome` builds Shiu et al.'s (2024) model of the whole fly brain from FlyWire. It reproduces their published runs, with a rate correlation of 0.999 and the motor neuron MN9 at 67.1 Hz against their 67.0 ± 6.6, at about 30 s per simulated second on 4 CPU cores. `FLYNN` (Wang and Chen 2026) trains a connectome as a recurrent rate network with one learned weight per synapse; against their PyTorch cell its activity and gradients agree within 1e-15.
+
+## RNeuralNet
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/messages-dark.webp">
+  <img alt="Messages travelling along the connections of a small RNeuralNet, each connection with its own delay" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/messages-light.webp" width="100%">
+</picture>
+
+`sparx.learn.RNeuralNet` rebuilds RNeuralNet-Research (2018), an early project of the author's, deterministically. Graded neurons sit on a random graph, each connection delivers its messages after its own delay, and a reward spreads backward by a softmax of activity. Compiled and run in a fixed order, the original C++ and sparx agree within 7.2e-7. On a delayed cue-order task, REINFORCE through the same network learns the task on four of five seeds, and the reward-diffusion rule never changes the network's choice. AGREL's update, a signed error sent back from the chosen output through the weights, learns it on the same four seeds; the same error spread by the original's shares does not.
+
+## Training on dew
+
+<picture>
+  <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-dark.svg">
+  <img alt="A Flax model and a sparx objective go to dew's trainer, which writes a run record that reloads, serves and exports" src="https://raw.githubusercontent.com/AshishKumar4/sparx/main/docs/assets/training-light.svg" width="100%">
+</picture>
+
+The `Trainer` from [dew](https://github.com/AshishKumar4/dew) runs sparx's objectives for classification, activity fitting, e-prop, predictive coding and RNeuralNet's rewards. A run's record names every class by import path, so `dew.pipeline("runs/shd", trust=("sparx",))` loads a trained network in a new process, and `sparx.serve.StreamServer` serves it to many streams at once. The [guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#training-on-dew) has a full SHD script.
+
+## Results
+
+| Task | Network | Test accuracy |
+| --- | --- | --- |
+| MNIST, rate-coded, 8 steps | 784-512-512 LIF | 97.5% after 2 epochs |
+| SHD, Hammouamri et al.'s recipe, 150 epochs, three seeds | 140-256-256 LIF with learned delays | 93.99 ± 0.29% at the last epoch (their code on the same GPU: 93.89 ± 0.26%) |
+| SHD, 140 channels | 140-128 ALIF, with and without learned delays | 74.6% and 64.5% |
+| Fashion-MNIST, Seely and Gould's headline cell | ReLU residual MLP, depth 32 | PC-ALM 75.1%, PC 62.2%, backpropagation 77.8% |
+| Pattern completion, Miconi et al.'s task | plastic recurrent network | 0.3% of bits wrong; 50.1% without fast weights |
+
+The SHD row is the full recipe beside the authors' code, both on an A100, three seeds each ([research/shd](https://github.com/AshishKumar4/sparx/blob/main/research/shd/README.md)). Both train on every training recording and score the test set after each epoch. The paper's 95.07 ± 0.24% (a 95% confidence interval over ten runs) is the best epoch on the test set, which chooses with the test set; here, as mean ± standard deviation over three seeds, their code's best epoch is 95.17 ± 0.61% and sparx's 94.96 ± 0.89%. With a tenth of the training set held out to choose the epoch, sparx scores 94.14 ± 0.98% on test. The other rows are short, untuned runs on a 4-core CPU; the [guide](https://github.com/AshishKumar4/sparx/blob/main/docs/guide.md#results-in-detail) gives their commands, times and comparisons.
+
+## Correctness
+
+Every model is checked against a reference: a float64 loop of its equations, the original authors' code, or NEST and Brian2. [docs/fidelity.md](https://github.com/AshishKumar4/sparx/blob/main/docs/fidelity.md) lists each model's reference, the check, the observed error and every known difference. `pytest -q` runs all of it on CPU in about 16 minutes.
+
+[`tools/make_figures.py`](https://github.com/AshishKumar4/sparx/blob/main/tools/make_figures.py) draws the banner and diagrams, and [`tools/make_clips.py`](https://github.com/AshishKumar4/sparx/blob/main/tools/make_clips.py) renders the clips. The spikes in the banner and the clips come from sparx runs.
+
+## License
+
+MIT

sparxml-0.1.0/pyproject.toml

@@ -0,0 +1,116 @@
+[build-system]
+requires = ["setuptools>=77", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "sparxml"
+dynamic = ["version"]
+description = "Spiking neural networks in JAX and Flax"
+readme = "README.md"
+requires-python = ">=3.12"
+authors = [
+  { name="Ashish Kumar Singh", email="ashishkmr472@gmail.com" }
+]
+license = "MIT"
+license-files = ["LICENSE"]
+keywords = ["jax", "flax", "spiking-neural-networks", "neuromorphic", "computational-neuroscience",
+            "connectomics", "dew"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Intended Audience :: Developers",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+    "Typing :: Typed",
+]
+dependencies = [
+    # Dew is the platform sparx trains, distributes, checkpoints and serves
+    # through (docs/design.md), and bounds jax, flax and optax to the versions
+    # it supports. Development and CI install the dew commit constraints.txt
+    # names.
+    "dewml>=0.1.0,<0.2",
+    "jax",
+    "flax",
+    "optax",
+    "numpy",
+]
+
+[project.urls]
+Homepage = "https://sparxml.dev"
+Repository = "https://github.com/AshishKumar4/sparx"
+Documentation = "https://sparxml.dev/docs/"
+Issues = "https://github.com/AshishKumar4/sparx/issues"
+Changelog = "https://github.com/AshishKumar4/sparx/blob/main/CHANGELOG.md"
+
+[project.optional-dependencies]
+# The accelerator builds of jax, as dew installs them.
+cuda12 = ["dewml[cuda12]>=0.1.0,<0.2"]
+cuda13 = ["dewml[cuda13]>=0.1.0,<0.2"]
+tpu = ["dewml[tpu]>=0.1.0,<0.2"]
+# sparx.datasets reads the published HDF5 files.
+datasets = ["h5py"]
+# sparx.graph.connectome reads connectome tables (CSV and parquet).
+connectome = ["pyarrow"]
+# sparx.nir exchanges networks through NIR.
+nir = ["nir"]
+test = ["pytest", "ruff==0.14.3", "h5py", "nir", "pyarrow"]
+
+[tool.setuptools.dynamic]
+version = { attr = "sparx.__version__" }
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+sparx = ["py.typed"]
+
+[tool.ruff]
+target-version = "py312"
+# dew's width (dew's pyproject.toml says why).
+line-length = 110
+
+[tool.ruff.lint]
+# dew's rules. Scripts use stdout as their interface; src/sparx/ruff.toml adds
+# T20 for the library, as dew's src/dew/ruff.toml does.
+select = [
+    "E", "F", "W", "I", "UP", "B", "C4", "PERF", "PIE", "RET", "SIM", "PGH",
+    "ERA", "ANN401", "FBT003", "PLW0120", "RUF", "C901", "PLR0912", "PLR0915",
+]
+allowed-confusables = ["×", "γ", "σ", "τ", "θ", "β", "α", "ρ"]
+
+[tool.ruff.lint.per-file-ignores]
+# dew's checker, verbatim at the commit sparx pins (CI checks it); its rule
+# walkers are as complex as dew's own limit of 25 allows.
+"tools/lint_slop.py" = ["C901"]
+
+[tool.ruff.lint.pylint]
+# dew's limits. Sparx keeps the stricter mccabe default of 10.
+max-branches = 25
+max-statements = 75
+
+[tool.ruff.lint.isort]
+known-first-party = ["sparx"]
+combine-as-imports = true
+split-on-trailing-comma = false
+
+[tool.pyright]
+# As dew's: sparx's own source and its checker, against the project venv. A
+# worktree has no .venv of its own, so the gate passes --pythonpath instead.
+include = ["src/sparx", "tools/lint_slop.py"]
+pythonVersion = "3.12"
+typeCheckingMode = "standard"
+venvPath = "."
+venv = ".venv"
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+pythonpath = ["src", "tests"]
+filterwarnings = [
+    'error::DeprecationWarning:sparx($|\.)',
+    'error::PendingDeprecationWarning:sparx($|\.)',
+    'error::FutureWarning:sparx($|\.)',
+]

sparxml-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

sparxml-0.1.0/src/sparx/__init__.py

@@ -0,0 +1,69 @@
+"""Sparx: spiking neural networks in JAX and Flax, trained, simulated and served through dew.
+
+Networks are Flax linen modules over time-major spike trains `[T, ...]`.
+
+Deep spiking networks:
+
+- `sparx.nn`: Flax layers over the neuron models, parallel spiking neurons and delayed synapses.
+- `sparx.models`: architectures built from them (`SEWResNet`, `SpikingMLP`).
+- `sparx.surrogate`: the spike and its surrogate gradients.
+- `sparx.encode`: the encoders that turn data into spike trains.
+- `sparx.losses` and `sparx.rates`: losses over time, firing-rate readouts and penalties.
+- `sparx.learn`: rules beyond backpropagation through time (e-prop, OTTT, EventProp, conversion).
+
+Circuits in physical units:
+
+- `sparx.dynamics`: neuron, synapse and plasticity models as pure JAX, the dimensionless family deep
+  networks train with and the physical one, and `run`, which scans any of them over time.
+- `sparx.graph`: populations and projections wired into a `Network`, `simulate`, and connectomes.
+- `sparx.spiketrains`: statistics of and distances between recorded spike trains.
+
+Around them:
+
+- `sparx.objectives`: the objectives that train spiking networks under dew's `Trainer`, with
+  `sparx.metrics` (their accuracy), `sparx.tasks` (the trained classifier `dew.pipeline` loads) and
+  `sparx.config` (`SNNRunConfig`, the run a recipe trains and `run.json` records).
+- `sparx.datasets`: spiking datasets as dew datasets (SHD).
+- `sparx.serve`: `StreamServer`, many streaming sessions in one batch.
+- `sparx.nir`: exchange through the Neuromorphic Intermediate Representation.
+
+`sparx.graph`, `sparx.learn`, `sparx.objectives`, `sparx.metrics`,
+`sparx.tasks`, `sparx.config`, `sparx.datasets`, `sparx.serve` and
+`sparx.nir` load on first access (`sparx.graph.Network` after `import
+sparx`). The graph, the objectives and the datasets import dew's trainer
+and data stack, about 0.9 s on a 4-core CPU, which a script that only
+trains a network in its own loop does not need.
+"""
+
+from __future__ import annotations
+
+import importlib
+from types import ModuleType
+from typing import TYPE_CHECKING
+
+from sparx import dynamics, encode, losses, models, nn, rates, spiketrains, surrogate
+from sparx.dynamics import run
+from sparx.rates import firing_rates, rate_penalty
+from sparx.surrogate import spike
+
+if TYPE_CHECKING:
+    from sparx import config, datasets, graph, learn, metrics, nir, objectives, serve, tasks
+
+__version__ = "0.1.0"
+
+_LAZY = ("config", "datasets", "graph", "learn", "metrics", "nir", "objectives", "serve", "tasks")
+
+__all__ = ["__version__", "config", "datasets", "dynamics", "encode", "firing_rates", "graph", "learn",
+           "losses", "metrics", "models", "nir", "nn", "objectives", "rate_penalty", "rates", "run", "serve",
+           "spike", "spiketrains", "surrogate", "tasks"]
+
+
+def __getattr__(name: str) -> ModuleType:
+    # PEP 562: importing the submodule binds it on the package, so this runs once per name.
+    if name in _LAZY:
+        return importlib.import_module(f"sparx.{name}")
+    raise AttributeError(f"module 'sparx' has no attribute {name!r}")
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_LAZY))