Fix/jit inside jit

[michalk8]

TODOs:

• create the "sharp bits"-like page

closes #312

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB

[codecov-commenter]

Codecov Report

Merging #335 (27bf935) into main (db689bd) will decrease coverage by 0.28%.
The diff coverage is 100.00%.

📣 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     #335      +/-   ##
==========================================
- Coverage   88.59%   88.32%   -0.28%     
==========================================
  Files          51       51              
  Lines        5479     5472       -7     
  Branches      824      824              
==========================================
- Hits         4854     4833      -21     
- Misses        509      522      +13     
- Partials      116      117       +1     

Impacted Files	Coverage Δ
src/ott/solvers/was_solver.py	88.09% <ø> (-7.26%)	⬇️
src/ott/solvers/linear/continuous_barycenter.py	93.75% <100.00%> (-0.10%)	⬇️
src/ott/solvers/linear/sinkhorn.py	96.83% <100.00%> (-0.02%)	⬇️
src/ott/solvers/linear/sinkhorn_lr.py	96.17% <100.00%> (-0.02%)	⬇️
src/ott/solvers/quadratic/gromov_wasserstein.py	89.31% <100.00%> (-3.87%)	⬇️
src/ott/solvers/quadratic/gw_barycenter.py	91.26% <100.00%> (-5.86%)	⬇️

[marcocuturi]

LGTM, we can either push or add the sharp bits page, as you prefer

[michalk8]

LGTM, we can either push or add the sharp bits page, as you prefer

Thanks, let's do the page later!