Add CI tests for upstream JAX

[michalk8]

closes #305

[codecov-commenter]

Codecov Report

Merging #313 (0015853) into main (ab9d89f) will increase coverage by 2.81%.
The diff coverage is 83.82%.

📣 This organization is not using Codecov’s GitHub App Integration. We recommend you install it so Codecov can continue to function properly for your repositories. Learn more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #313      +/-   ##
==========================================
+ Coverage   85.95%   88.76%   +2.81%     
==========================================
  Files          52       52              
  Lines        5590     5592       +2     
  Branches      853      582     -271     
==========================================
+ Hits         4805     4964     +159     
+ Misses        662      497     -165     
- Partials      123      131       +8     

Impacted Files	Coverage Δ
src/ott/problems/linear/potentials.py	62.42% <42.85%> (-1.40%)	⬇️
src/ott/tools/plot.py	19.82% <55.55%> (+<0.01%)	⬆️
src/ott/problems/quadratic/quadratic_problem.py	88.76% <66.66%> (-0.57%)	⬇️
src/ott/geometry/geometry.py	92.40% <92.59%> (+1.23%)	⬆️
src/ott/geometry/epsilon_scheduler.py	94.44% <100.00%> (+5.55%)	⬆️
src/ott/geometry/graph.py	91.87% <100.00%> (+60.00%)	⬆️
src/ott/geometry/grid.py	91.60% <100.00%> (-0.07%)	⬇️
src/ott/geometry/low_rank.py	96.37% <100.00%> (-0.03%)	⬇️
src/ott/geometry/pointcloud.py	86.04% <100.00%> (-0.09%)	⬇️
src/ott/problems/nn/dataset.py	94.73% <0.00%> (-0.14%)	⬇️
... and 5 more