Make using JAX (or any accelerator) an Instanced Python class (and toggle)#509

Open

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

+1,496-1,010

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

Commits on Jun 22, 2023

Commits on Jun 26, 2023

Commits on Jun 27, 2023

Fix the docs, may conflict with main due to choderalab#510 but I can fix that too.
Lnaden
committed