scvi-tools
scvi-tools copied to clipboard
Implementation of PeakVI in Jax
Reimplementation of PeakVI in Jax. This is meant to both see if we can get a similar speed up in PeakVI, as well as understanding the limitations of the JaxModuleWrapper abstraction.