Timings and Code for Spiking Neural Networks with JAX
I've been encouraged to flesh out my earlier posts about JAX to support 27DaysOfJAX . I've written simulations of a Leaky Integrate and Fire Neuron in *Plowman's* (pure) Python, Python + numpy, and Python + JAX. Here's a plot of a 2000-step simulation for a single neuron: Plot for a single neuron The speedups using Python, Jax and the JAX jit compiler are dramatic. Pure Python can simulate a single step for a single neuron in roughly 0.25 µs. so 1,000,000 neurons would take about 0.25 seconds. numpy can simulate a single step for 1,000,000 neurons in 13.7 ms . Python, JAX + JAX's jit compilation can simulate a single step for 1,000,000 neurons in 75 µs . Here's the core code for each version. # Pure Python def step(v, tr, injected_current): spiking = False if tr > 0: next_v = reset_voltage tr = tr - 1 elif v > threshold: next_v = reset_voltage tr = int(refactory_period / dt) spiking = True else...