jaxon reimplements the canonical myelinated (MRG, Sweeney) and unmyelinated (Sundt, Rattay) axon models in JAX/Jaxley, replacing NEURON's serial, non-differentiable forward solver with a vmap-batched, autodiff-capable coupled intracellular/periaxonal double-cable solver .
Validation evidence is strong for what it tests: 943 activation-threshold configurations agree with PyFibers-wrapped NEURON within 1% in 99.6% of cases; conduction velocity matches to machine precision; and the block-Thomas sweep agrees with a dense LU reference to 5.6Γ10β»ΒΉΒΉ mV . Conduction velocity agreement to machine precision is consistent with both codes integrating identical equations; it validates the implementation, not the biology.
Speedup claim is narrower than headline suggests. The geometric-mean ~820Γ figure applies at N=100,000 fibers on a single GPU; the paper does not report per-N scaling curves in the text, hardware details, or a comparison against AxonML surrogates β a notable omission since AxonML is cited as the prior GPU alternative .
The genuinely new capability is differentiable stimulus optimization through the exact cable equation (per-contact amplitudes, arbitrary waveforms, electrode position) rather than grid search β a real advance over both NEURON and learned surrogates, because the gradient is of the biophysics itself . However: (1) no optimized stimulus is directly compared against grid-search or AxonML in this paper β the authors defer that application to a companion preprint whose DOI is not yet assigned; (2) no author-reported limitations section exists (inferred limitations are ours, not the authors'); (3) thresholds use a peak-Na m-gate activation proxy rather than a direct spike criterion, which should be verified against edge cases near threshold; (4) reproducibility is excellent on paper (MIT license, Zenodo code + CC-BY-4.0 data, REPRODUCE.md) but has not yet been independently replicated.
What would change this assessment: an independent replication falling below 99.6% agreement, a multi-GPU/other-GPU benchmark failing to reproduce the speedup, or a head-to-head showing AxonML-style surrogates achieving comparable optimization quality at lower cost .
Know what changed, what holds up, and what remains uncertain. Every Friday. No ads.