bioRxiv · 10.64898/2026.07.30.741846
jaxon: a differentiable, GPU-native simulator for peripheral-nerve fiber models
Abstract
jaxon is an open-source, fully differentiable and GPU-native reimplementation of the canonical peripheral-nerve fiber models in JAX/Jaxley: the myelinated McIntyre-Richardson-Grill (MRG) and Sweeney axons and the un-myelinated Sundt and Rattay C-fibers. It reproduces NEURONs extracellular mechanism through a custom backward-Euler coupled intracellular/periaxonal double-cable solver, agreeing with PyFibers-wrapped NEURON on 99.6% of 943 activation-threshold configurations within 1% and matching conduction velocity to machine precision. Because the entire forward model is expressed in JAX, it is both vectorized--simulating whole fiber populations in parallel and reaching a geometric-mean [~]820x speedup at N = 100,000 fibers on a single GPU--and differentiable, so extracellular-stimulation parameters (per-contact amplitudes, waveform shape, and electrode position) can be optimized directly through the cable equation rather than grid-searched. jaxon slots into existing peripheral-nerve modeling pipelines as a gradient-enabled, population-scale replacement for the NEURON forward solver.
Explore related subjects
Keep this discovery
Explore connections, maps & timelines
Lung, D., Haberbusch, M.. 2026-07-31. jaxon: a differentiable, GPU-native simulator for peripheral-nerve fiber models. https://doi.org/10.64898/2026.07.30.741846
Cite the original work for its findings. Save a collection to share your selection of sources.