refactor: Dropping Jax from scvi-tools

[ori-kron-wis]

No description provided.

[codecov]

Codecov Report

❌ Patch coverage is 50.00000% with 6 lines in your changes missing coverage. Please review.
✅ Project coverage is 84.17%. Comparing base (4b1bb80) to head (41154f3).

Files with missing lines	Patch %	Lines
src/scvi/external/mrvi/_model.py	33.33%	4 Missing ⚠️
src/scvi/model/_utils.py	33.33%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #3786      +/-   ##
==========================================
- Coverage   88.22%   84.17%   -4.05%     
==========================================
  Files         229      229              
  Lines       22634    22197     -437     
==========================================
- Hits        19969    18685    -1284     
- Misses       2665     3512     +847     

Flag	Coverage Δ
integration	75.15% <50.00%> (+3.68%)	⬆️
nonjax	?

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
src/scvi/__init__.py	100.00% <ø> (ø)
src/scvi/_settings.py	90.26% <ø> (+3.24%)	⬆️
src/scvi/_types.py	100.00% <100.00%> (+7.14%)	⬆️
src/scvi/distributions/__init__.py	100.00% <ø> (+9.09%)	⬆️
src/scvi/distributions/_negative_binomial.py	78.73% <ø> (-1.51%)	⬇️
src/scvi/external/__init__.py	100.00% <ø> (+14.70%)	⬆️
src/scvi/external/mrvi_torch/_model.py	88.30% <ø> (ø)
src/scvi/model/__init__.py	95.65% <ø> (-0.78%)	⬇️
src/scvi/model/base/__init__.py	94.73% <ø> (-1.10%)	⬇️
src/scvi/module/__init__.py	95.23% <ø> (-0.92%)	⬇️
... and 11 more

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 21296cc828

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[ori-kron-wis]

We will need to transform Tangram to Torch before we remove Jax. Afterwards it will also move to scviva-tools @canergen