Add opt-in memory-efficient CSR aggregation for graded synapses

[TorgemanTarak]

FlyVis's scatter aggregation retains edge-sized activity tensors during recurrent training. This adds an optional Network(backend="sparse_trainable") backend for the default graded-synapse dynamics, reducing activation memory while preserving the existing outputs, first-order gradients, parameter sharing, and checkpoint keys. Scatter remains the default.

The implementation coalesces parallel edges, caches both CSR layouts, and computes sparse weight and activity gradients without retaining batch-by-edge tensors. It supports CPU/CUDA float32 and float64 and the built-in pointwise activations. Custom dynamics, mixed precision, higher-order gradients, and compilation are outside this patch's supported scope.

On the supplied visual-system graph (45,669 nodes, 1,513,231 edges), an RTX 5090 benchmark at batch 16 and 32 steps reduced peak allocated memory from 3.641 to 0.886 GiB (75.7%) and median forward/backward time from 48.02 to 30.95 ms (1.55x). Batch 4 saved memory but was slower. These are synthetic forward/backward measurements, excluding optimizer updates; they do not establish task accuracy or time to convergence. Reproduction instructions and raw measurements are included in benchmarks/.

Validation: 198 tests passed and 12 skipped across new sparse tests and existing network, dynamics, and initialization suites. Coverage includes output/gradient parity, finite-difference checks, shared parameters, duplicate edges, initial-state gradients, optimizer-step parity, checkpoint compatibility, and CPU/CUDA execution. Ruff and whitespace checks passed. The full pretrained-asset suite was not run. An isolated Windows datamate cache-creation workaround was required and is documented with the benchmark; it is not part of this source patch.

[Copilot]

Copilot review overview

🔵 Needs a closer look

The custom autograd and CUDA sparse operations warrant final human review despite comprehensive parity testing.

Review effort: Balanced
Findings: None

What changed in this PR

Adds an opt-in CSR backend for memory-efficient graded-synapse training while retaining scatter as the default.

Changes:

• Implements custom sparse forward and first-order backward aggregation.
• Integrates backend selection, validation, configuration, and parity tests.
• Documents limitations and provides reproducible RTX 5090 benchmarks.

File	Description
flyvis/​network/​sparse.py	Implements CSR aggregation and gradients.
flyvis/​network/​network.py	Integrates backend selection and sparse weights.
flyvis/​network/​dynamics.py	Supports node-level current aggregation.
tests/​test_sparse_current.py	Tests parity, gradients, edge cases, and checkpoints.
docs/​docs/​reference/​network.md	Documents usage and limitations.
benchmarks/​sparse_current.py	Adds the benchmark runner.
benchmarks/​rtx5090-sparse-training.md	Summarizes benchmark methodology and results.
benchmarks/​rtx5090-sparse-training.jsonl	Records raw benchmark measurements.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.