Feature/batched vmap

[michalk8]

TODOs:

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

closes #504

[codecov]

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 ⚠️
src/ott/utils.py	98.05%	1 Missing and 1 partial ⚠️

Additional details and impacted files

@@            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	Coverage Δ
src/ott/geometry/costs.py	97.18% <100.00%> (-0.04%)	⬇️
src/ott/geometry/distrib_costs.py	100.00% <100.00%> (ø)
src/ott/geometry/geodesic.py	94.39% <100.00%> (+0.05%)	⬆️
src/ott/geometry/geometry.py	94.46% <100.00%> (+1.21%)	⬆️
src/ott/geometry/graph.py	95.83% <100.00%> (+0.04%)	⬆️
src/ott/geometry/grid.py	96.24% <100.00%> (+0.02%)	⬆️
src/ott/geometry/low_rank.py	96.73% <100.00%> (-0.31%)	⬇️
src/ott/neural/methods/monge_gap.py	91.20% <ø> (ø)
src/ott/problems/linear/potentials.py	91.50% <ø> (ø)
src/ott/problems/quadratic/gw_barycenter.py	90.00% <ø> (ø)
... and 7 more

[marcocuturi]

LGTM!

[marcocuturi]

true!