update nb (#403)

* update nb

* ref

* ref

* no need to add ref!

* pydocs

* bump flax

* comments + orbax fix?

* flax?

* Pin `flax` version in the CI

* Try `--no-deps`

* Fix some math rendering

---------

Co-authored-by: Michal Klein <46717574+michalk8@users.noreply.github.com>

.github/workflows/tests.yml

@@ -56,11 +56,12 @@ jobs:
     - name: Install dependencies
       # `jax[cuda]<0.4` because of Docker issues: https://github.com/google/jax/issues/13758
       # `chex<0.1.7` because it requires `jax>=0.4.6`
-      # `flax<0.6.5` because it requires `jax>=0.4.2`
+      # `flax<0.6.5` because it requires `jax>=0.4.2`, --no-deps because of `orbax`
       run: |
         python3 -m pip install --upgrade pip
         python3 -m pip install -e".[test]"
-        python3 -m pip install "flax<0.6.5" "chex<0.1.7"
+        python3 -m pip install "orbax-checkpoint" "orbax-export" "chex<0.1.7"
+        python3 -m pip install --no-deps "flax<0.6.5"
         python3 -m pip install "jax[cuda]<0.4" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
 
     - name: Nvidia SMI

docs/tutorials/notebooks/MetaOT.ipynb

@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "id": "7eb605bd",
    "metadata": {},
@@ -535,6 +534,7 @@
     "We lastly compare how much the initializers help\n",
     "{class}`~ott.solvers.linear.sinkhorn.Sinkhorn` converge on these problems, \n",
     "measured by the marginal error:\n",
+    "\n",
     "$${\\rm err}(f,g; \\alpha, \\beta, c) := \\|P1_m-a\\|_1 + \\|P^\\top1_n-b\\|_1$$"
    ]
   },

docs/tutorials/notebooks/OTT_&_POT.ipynb

@@ -74,17 +74,19 @@
     "$$\\mu=\\sum_{i=1}^n a_i\\delta_{x_i}, \\nu =\\sum_{j=1}^n b_j\\delta_{y_j},$$\n",
     "\n",
     "to define the OT problem in its primal form,\n",
+    "\n",
     "$$\\min_{P \\in U(a,b)} \\langle C, P \\rangle - \\varepsilon H(P).$$\n",
     "\n",
     "where $U(a,b):=\\{P \\in \\mathbf{R}_+^{n\\times n}, P\\mathbf{1}_{n}=b, P^T\\mathbf{1}_n=b\\}$, and $C = [ \\|x_i - y_j \\|^2 ]_{i,j}\\in \\mathbf{R}_+^{n\\times n}$.\n",
     "\n",
     "That problem is equivalent to the following dual form,\n",
+    "\n",
     "$$\\max_{f, g} \\langle a, f \\rangle + \\langle b, g \\rangle - \\varepsilon \\langle e^{f/\\varepsilon},Ke^{g/\\varepsilon} \\rangle.$$\n",
     "\n",
     "These two problems are solved by `OTT` and `POT` using the {class}`~ott.solvers.linear.sinkhorn.Sinkhorn` iterations with a simple initialization for $u$, and subsequent updates $v \\leftarrow a / K^Tu, u \\leftarrow b / Kv$,  where $K:=e^{-C/\\varepsilon}$.\n",
     "\n",
-    "Upon convergence to fixed points $u^*, v^*$, one has $$P^*=D(u^*)KD(v^*)$$ or, alternatively, \n",
-    "$$f^*, g^* = \\varepsilon \\log(u^*), \\varepsilon\\log(v^*)$$"
+    "Upon convergence to fixed points $u^*, v^*$, one has $P^*=D(u^*)KD(v^*)$ or, alternatively, \n",
+    "$f^*, g^* = \\varepsilon \\log(u^*), \\varepsilon\\log(v^*)$."
    ]
   },
   {

pyproject.toml

@@ -18,7 +18,7 @@ dependencies = [
     "jaxopt>=0.5.5",
     # numba/numpy compatibility issue in JAXOPT.
     "numpy>=1.18.4, <1.25.0",
-    "flax>=0.5.2",
+    "flax>=0.6.6",
     "optax>=0.1.1",
     "lineax>=0.0.1; python_version >= '3.9'"
 ]

src/ott/geometry/costs.py

@@ -315,13 +315,14 @@ class RegTICost(TICost, abc.ABC):
   r"""Base class for regularized translation-invariant costs.
 
   .. math::
-    \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} reg\left(\cdot\right)
+    \frac{1}{2} \|\cdot\|_2^2 + \text{scaling_reg} reg\left(matrix \cdot\right)
 
   where :func:`reg` is the regularization function.
 
   Args:
     scaling_reg: Strength of the :meth:`regularization <reg>`.
-    matrix: :math:`p \times d` projection matrix with **orthogonal rows**.
+    matrix: :math:`p \times d` projection matrix in the Stiefel manifold,
+      namely with **orthonormalized rows**.
     orthogonal: Whether to regularize in the orthogonal complement
       to promote displacements in the span of ``matrix``.
   """