test: removing jax from test-suite

[ori-kron-wis]

The idea here is to simulate the case where scvi-tools is being installed w/o jax, as it is starting 1.4.0, and then check all the tests, making sure everything works (beside the JAX-related tests)

[codecov]

Codecov Report

❌ Patch coverage is 80.00000% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 74.15%. Comparing base (6fc1538) to head (14b6ddc).
⚠️ Report is 2 commits behind head on main.

Files with missing lines	Patch %	Lines
src/scvi/distributions/_negative_binomial.py	80.00%	2 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (6fc1538) and HEAD (14b6ddc). Click for more details.

HEAD has 63 uploads less than BASE

Flag	BASE (6fc1538)	HEAD (14b6ddc)
	67	4

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #3558       +/-   ##
===========================================
- Coverage   87.30%   74.15%   -13.15%     
===========================================
  Files         224      224               
  Lines       21011    21013        +2     
===========================================
- Hits        18344    15583     -2761     
- Misses       2667     5430     +2763     

Files with missing lines	Coverage Δ
src/scvi/distributions/_negative_binomial.py	83.39% <80.00%> (-1.07%)	⬇️

... and 61 files with indirect coverage changes

🚀 New features to boost your workflow:

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