Make using JAX (or any accelerator) a toggle#508

Closed

Lnaden wants to merge 8 commits intochoderalab:masterfrom Lnaden:jax-toggle

+1,537-1,008

Commits on Jun 16, 2023

First attempted implementation of the API for toggling JAX. Needs lots of work still, including how to delay JIT
Lnaden
committed
Make using JAX a toggle
Lnaden
committed

Commits on Jun 20, 2023

First take on implementing a much more pythonic jax-toggle system.
Lnaden
committed
Next iterations, trying to JIT without causing all kinds of pain of recompile. Seems to run much slower in tests right now.
Lnaden
committed