search

scverse / scvi-tools

Deep probabilistic analysis of single-cell and spatial omics data
http://scvi-tools.org/
BSD 3-Clause "New" or "Revised" License
1.21k stars 344 forks source link

scvi-tools Jax support #1718

Open martinkim0 opened 2 years ago

martinkim0 commented 2 years ago

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

martinkim0 commented 2 years ago

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 commented 2 years ago

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 commented 2 years ago

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.