Feature/batched vmap

[michalk8]

TODOs:

• [x] negative in/out axes
• [x] clean-up the pointcloud.py (re-order methods/properties, etc.)
• [x] update LRCGeometry
• [x] test the batched_vmap
• [x] docs

closes #504

[codecov[bot]]

Codecov Report

Attention: Patch coverage is 97.55245% with 7 lines in your changes missing coverage. Please review.

Project coverage is 88.13%. Comparing base (706cef7) to head (e390e64). Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/ott/geometry/pointcloud.py	95.79%	3 Missing and 2 partials :warning:
src/ott/utils.py	98.05%	1 Missing and 1 partial :warning:

Additional details and impacted files

[](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax) ```diff @@ Coverage Diff @@ ## main #588 +/- ## ========================================== + Coverage 88.00% 88.13% +0.12% ========================================== Files 73 73 Lines 7820 7768 -52 Branches 567 556 -11 ========================================== - Hits 6882 6846 -36 + Misses 789 779 -10 + Partials 149 143 -6 ``` | [Files with missing lines](https://app.codecov.io/gh/ott-jax/ott/pull/588?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax) | Coverage Δ | | |---|---|---| | [src/ott/geometry/costs.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Fcosts.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9jb3N0cy5weQ==) | `97.18% <100.00%> (-0.04%)` | :arrow_down: | | [src/ott/geometry/distrib\_costs.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Fdistrib_costs.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9kaXN0cmliX2Nvc3RzLnB5) | `100.00% <100.00%> (ø)` | | | [src/ott/geometry/geodesic.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Fgeodesic.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9nZW9kZXNpYy5weQ==) | `94.39% <100.00%> (+0.05%)` | :arrow_up: | | [src/ott/geometry/geometry.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Fgeometry.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9nZW9tZXRyeS5weQ==) | `94.46% <100.00%> (+1.21%)` | :arrow_up: | | [src/ott/geometry/graph.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Fgraph.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9ncmFwaC5weQ==) | `95.83% <100.00%> (+0.04%)` | :arrow_up: | | [src/ott/geometry/grid.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Fgrid.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9ncmlkLnB5) | `96.24% <100.00%> (+0.02%)` | :arrow_up: | | [src/ott/geometry/low\_rank.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fgeometry%2Flow_rank.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9nZW9tZXRyeS9sb3dfcmFuay5weQ==) | `96.73% <100.00%> (-0.31%)` | :arrow_down: | | [src/ott/neural/methods/monge\_gap.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fneural%2Fmethods%2Fmonge_gap.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9uZXVyYWwvbWV0aG9kcy9tb25nZV9nYXAucHk=) | `91.20% <ø> (ø)` | | | [src/ott/problems/linear/potentials.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fproblems%2Flinear%2Fpotentials.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9wcm9ibGVtcy9saW5lYXIvcG90ZW50aWFscy5weQ==) | `91.50% <ø> (ø)` | | | [src/ott/problems/quadratic/gw\_barycenter.py](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree&filepath=src%2Fott%2Fproblems%2Fquadratic%2Fgw_barycenter.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax#diff-c3JjL290dC9wcm9ibGVtcy9xdWFkcmF0aWMvZ3dfYmFyeWNlbnRlci5weQ==) | `90.00% <ø> (ø)` | | | ... and [7 more](https://app.codecov.io/gh/ott-jax/ott/pull/588?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=ott-jax) | |

[review-notebook-app[bot]]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB