jax
brainunit
brainstate
braintools

[cpu]
jaxlib

[cuda12]
jaxlib[cuda12]

[cuda13]
jaxlib[cuda13]

[testing]
pytest

[tpu]
jaxlib[tpu]
