scvi-tools Jax support

[martinkim0]

This issue will serve as a parent thread for all issues related to supporting Jax-backend for scvi-tools models.

• FCLayers analog Jax analog for FCLayers #1620
• PeakVI implementation Implementation of PeakVI in Jax #1578
• Granular device management Granular device management for Jax #1553
• Default Jax seed Default Jax seed to scvi.settings.seed #1584

[martinkim0]

I think in the long term, it would be useful if our high-level API is framework-agnostic. Something like the following:

model = scvi.model.SCVI(adata, backend="jax")

or

model = scvi.model.SCVI(adata, backend="torch")

[adamgayoso]

I think in the long term, it would be useful if our high-level API is framework-agnostic

We thought about this but it requires matching features exactly, which might be hard. But we should continue to evaluate this option.

[martinkim0]

We thought about this but it requires matching features exactly, which might be hard. But we should continue to evaluate this option.

Oh I see, yeah in that case we can revisit this when more code is implemented.