ott-jax (pypi): Optimal Transport Tools in JAX.
Keywords: optimal transport|gromov wasserstein|sinkhorn|low-rank sinkhorn|sinkhorn divergences|wasserstein|wasserstein barycenter|jax|autodiff|implicit differentiation|automatic-differentiation|optimal-transport