jaxleyverse / jaxley

Differentiable neuron simulations with biophysical detail on CPU, GPU, or TPU.
https://jaxley.readthedocs.io
Apache License 2.0
46 stars 7 forks source link

Make types consistent across Module attributes #454

Open jnsbck opened 1 week ago

jnsbck commented 1 week ago

Some attrs like cumsum_nbranches have different types in different Modules. Also we should make sure that the attributes are ideally numpy arrays not jax arrays.

fyi @michaeldeistler