Backport PR #3558 on branch 1.4.x (test: removing jax from test-suite)

[meeseeksmachine]

Backport PR #3558: test: removing jax from test-suite

[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 (e051bba) to head (0271ff2).
⚠️ Report is 1 commits behind head on 1.4.x.

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

Additional details and impacted files

@@           Coverage Diff           @@
##            1.4.x    #3559   +/-   ##
=======================================
  Coverage   74.15%   74.15%           
=======================================
  Files         224      224           
  Lines       21011    21013    +2     
=======================================
+ Hits        15581    15583    +2     
  Misses       5430     5430           

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

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:

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