ott-jax · michalk8 · Sep 13, 2023 · Sep 13, 2023 · Sep 13, 2023 · Sep 13, 2023
diff --git a/docs/index.rst b/docs/index.rst
@@ -103,7 +103,7 @@ Packages
     :maxdepth: 1
     :caption: Examples
 
-    Getting Started <tutorials/notebooks/basic_ot_between_datasets>
+    Getting Started <tutorials/basic_ot_between_datasets>
     tutorials/index
 
 .. toctree::

diff --git a/docs/tutorials/notebooks/GWLRSinkhorn.ipynb → docs/tutorials/GWLRSinkhorn.ipynb b/docs/tutorials/notebooks/GWLRSinkhorn.ipynb → docs/tutorials/GWLRSinkhorn.ipynb
diff --git a/docs/tutorials/notebooks/Hessians.ipynb → docs/tutorials/Hessians.ipynb b/docs/tutorials/notebooks/Hessians.ipynb → docs/tutorials/Hessians.ipynb
diff --git a/docs/tutorials/notebooks/LRSinkhorn.ipynb → docs/tutorials/LRSinkhorn.ipynb b/docs/tutorials/notebooks/LRSinkhorn.ipynb → docs/tutorials/LRSinkhorn.ipynb
diff --git a/docs/tutorials/notebooks/MetaOT.ipynb → docs/tutorials/MetaOT.ipynb b/docs/tutorials/notebooks/MetaOT.ipynb → docs/tutorials/MetaOT.ipynb
diff --git a/docs/tutorials/notebooks/Monge_Gap.ipynb → docs/tutorials/Monge_Gap.ipynb b/docs/tutorials/notebooks/Monge_Gap.ipynb → docs/tutorials/Monge_Gap.ipynb
@@ -1,7 +1,6 @@
 {
  "cells": [
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -22,7 +21,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -38,7 +37,6 @@
     "\n",
     "from matplotlib import pyplot as plt\n",
     "\n",
-    "%pip install -e ../../../\n",
     "from ott.geometry import costs, pointcloud\n",
     "from ott.problems.nn import dataset\n",
     "from ott.solvers.linear import acceleration\n",
@@ -47,17 +45,6 @@
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": 3,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "%load_ext autoreload\n",
-    "%autoreload 2"
-   ]
-  },
-  {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -83,7 +70,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -199,7 +185,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -293,7 +278,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -340,7 +324,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -464,7 +447,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -531,7 +513,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -575,7 +556,6 @@
    ]
   },
   {
-   "attachments": {},
    "cell_type": "markdown",
    "metadata": {},
    "source": [
@@ -649,9 +629,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "ott",
    "language": "python",
-   "name": "python3"
+   "name": "ott"
   },
   "language_info": {
    "codemirror_mode": {
@@ -663,7 +643,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,

diff --git a/docs/tutorials/notebooks/OTT_&_POT.ipynb → docs/tutorials/OTT_&_POT.ipynb b/docs/tutorials/notebooks/OTT_&_POT.ipynb → docs/tutorials/OTT_&_POT.ipynb
diff --git a/docs/tutorials/notebooks/One_Sinkhorn.ipynb → docs/tutorials/One_Sinkhorn.ipynb b/docs/tutorials/notebooks/One_Sinkhorn.ipynb → docs/tutorials/One_Sinkhorn.ipynb
diff --git a/...ials/notebooks/Sinkhorn_Barycenters.ipynb → docs/tutorials/Sinkhorn_Barycenters.ipynb b/...ials/notebooks/Sinkhorn_Barycenters.ipynb → docs/tutorials/Sinkhorn_Barycenters.ipynb
diff --git a/...rials/notebooks/application_biology.ipynb → docs/tutorials/application_biology.ipynb b/...rials/notebooks/application_biology.ipynb → docs/tutorials/application_biology.ipynb
diff --git a/...notebooks/basic_ot_between_datasets.ipynb → ...tutorials/basic_ot_between_datasets.ipynb b/...notebooks/basic_ot_between_datasets.ipynb → ...tutorials/basic_ot_between_datasets.ipynb
diff --git a/docs/tutorials/notebooks/gmm_pair_demo.ipynb → docs/tutorials/gmm_pair_demo.ipynb b/docs/tutorials/notebooks/gmm_pair_demo.ipynb → docs/tutorials/gmm_pair_demo.ipynb
diff --git a/docs/tutorials/gromov_wasserstein.ipynb b/docs/tutorials/gromov_wasserstein.ipynb
diff --git a/...books/gromov_wasserstein_multiomics.ipynb → ...rials/gromov_wasserstein_multiomics.ipynb b/...books/gromov_wasserstein_multiomics.ipynb → ...rials/gromov_wasserstein_multiomics.ipynb
diff --git a/docs/tutorials/notebooks/icnn_inits.ipynb → docs/tutorials/icnn_inits.ipynb b/docs/tutorials/notebooks/icnn_inits.ipynb → docs/tutorials/icnn_inits.ipynb
diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst
@@ -6,54 +6,54 @@ Geometry
 .. toctree::
     :maxdepth: 1
 
-    notebooks/introduction_grid
+    introduction_grid
 
 Linear Optimal Transport
 ------------------------
 .. toctree::
     :maxdepth: 1
 
-    notebooks/point_clouds
-    notebooks/One_Sinkhorn
-    notebooks/OTT_&_POT
-    notebooks/Hessians
-    notebooks/LRSinkhorn
-    notebooks/sinkhorn_divergence_gradient_flow
-    notebooks/sparse_monge_displacements
+    point_clouds
+    One_Sinkhorn
+    OTT_&_POT
+    Hessians
+    LRSinkhorn
+    sinkhorn_divergence_gradient_flow
+    sparse_monge_displacements
 
 Barycenters
 ^^^^^^^^^^^
 .. toctree::
     :maxdepth: 1
 
-    notebooks/Sinkhorn_Barycenters
-    notebooks/gmm_pair_demo
-    notebooks/wasserstein_barycenters_gmms
+    Sinkhorn_Barycenters
+    gmm_pair_demo
+    wasserstein_barycenters_gmms
 
 Miscellaneous
 ^^^^^^^^^^^^^
 .. toctree::
     :maxdepth: 1
 
-    notebooks/tracking_progress
-    notebooks/soft_sort
-    notebooks/application_biology
+    tracking_progress
+    soft_sort
+    application_biology
 
 Quadratic Optimal Transport
 ---------------------------
 .. toctree::
     :maxdepth: 1
 
-    notebooks/gromov_wasserstein
-    notebooks/GWLRSinkhorn
-    notebooks/gromov_wasserstein_multiomics
+    gromov_wasserstein
+    GWLRSinkhorn
+    gromov_wasserstein_multiomics
 
 Neural Optimal Transport
 ------------------------
 .. toctree::
     :maxdepth: 1
 
-    notebooks/neural_dual
-    notebooks/icnn_inits
-    notebooks/MetaOT
-    notebooks/Monge_Gap
+    neural_dual
+    icnn_inits
+    MetaOT
+    Monge_Gap
diff --git a/...torials/notebooks/introduction_grid.ipynb → docs/tutorials/introduction_grid.ipynb b/...torials/notebooks/introduction_grid.ipynb → docs/tutorials/introduction_grid.ipynb
diff --git a/docs/tutorials/notebooks/neural_dual.ipynb → docs/tutorials/neural_dual.ipynb b/docs/tutorials/notebooks/neural_dual.ipynb → docs/tutorials/neural_dual.ipynb
diff --git a/docs/tutorials/notebooks/gromov_wasserstein.ipynb b/docs/tutorials/notebooks/gromov_wasserstein.ipynb
diff --git a/docs/tutorials/notebooks/point_clouds.ipynb → docs/tutorials/point_clouds.ipynb b/docs/tutorials/notebooks/point_clouds.ipynb → docs/tutorials/point_clouds.ipynb
diff --git a/...s/sinkhorn_divergence_gradient_flow.ipynb → ...s/sinkhorn_divergence_gradient_flow.ipynb b/...s/sinkhorn_divergence_gradient_flow.ipynb → ...s/sinkhorn_divergence_gradient_flow.ipynb
diff --git a/docs/tutorials/notebooks/soft_sort.ipynb → docs/tutorials/soft_sort.ipynb b/docs/tutorials/notebooks/soft_sort.ipynb → docs/tutorials/soft_sort.ipynb
diff --git a/...otebooks/sparse_monge_displacements.ipynb → ...utorials/sparse_monge_displacements.ipynb b/...otebooks/sparse_monge_displacements.ipynb → ...utorials/sparse_monge_displacements.ipynb
diff --git a/...torials/notebooks/tracking_progress.ipynb → docs/tutorials/tracking_progress.ipynb b/...torials/notebooks/tracking_progress.ipynb → docs/tutorials/tracking_progress.ipynb
diff --git a/...ebooks/wasserstein_barycenters_gmms.ipynb → ...orials/wasserstein_barycenters_gmms.ipynb b/...ebooks/wasserstein_barycenters_gmms.ipynb → ...orials/wasserstein_barycenters_gmms.ipynb
diff --git a/pyproject.toml b/pyproject.toml
@@ -183,7 +183,10 @@ legacy_tox_ini = """
     skip_missing_interpreters = true
 
     [testenv]
-    extras = test,neural
+    extras =
+        test
+        # https://github.com/google/flax/issues/3329
+        py{3.9,3.10,3.11},py3.9-jax-default: neural
     pass_env = CUDA_*,PYTEST_*,CI
     commands_pre =
         gpu: python -I -m pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
@@ -299,7 +302,8 @@ select = [
 unfixable = ["B", "UP", "C4", "BLE", "T20", "RET"]
 target-version = "py38"
 [tool.ruff.per-file-ignores]
-"tests/*" = ["D", "PT004"]  # TODO(michalk8): remove `self.initialize` in `tests/`
+# TODO(michalk8): PO004 - remove `self.initialize`
+"tests/*" = ["D", "PT004", "E402"]
 "*/__init__.py" = ["F401"]
 "docs/*" = ["D"]
 "src/ott/types.py" = ["D102"]

diff --git a/src/ott/problems/linear/barycenter_problem.py b/src/ott/problems/linear/barycenter_problem.py
@@ -18,7 +18,7 @@
 
 from ott.geometry import costs, geometry, segment
 
-__all__ = ["FreeBarycenterProblem"]
+__all__ = ["FreeBarycenterProblem", "FixedBarycenterProblem"]
 
 
 @jax.tree_util.register_pytree_node_class
@@ -44,8 +44,8 @@ class FreeBarycenterProblem:
       Only used when ``y`` is not already segmented. When passing
       ``segment_ids``, 2 arguments must be specified for jitting to work:
 
-        - ``num_segments`` - the total number of measures.
-        - ``max_measure_size`` -  maximum of support sizes of these measures.
+      - ``num_segments`` - the total number of measures.
+      - ``max_measure_size`` -  maximum of support sizes of these measures.
   """
 
   def __init__(
@@ -158,7 +158,7 @@ class FixedBarycenterProblem:
     a: batch of histograms of shape ``[batch, num_a]`` where ``num_a`` matches
       the first value of the :attr:`~ott.geometry.Geometry.shape` attribute of
       ``geom``.
-    weights: ``[batch,]`` positive weights summing to :math`1`. Uniform by
+    weights: ``[batch,]`` positive weights summing to :math:`1`. Uniform by
       default.
   """
 

diff --git a/src/ott/solvers/linear/sinkhorn.py b/src/ott/solvers/linear/sinkhorn.py
@@ -318,7 +318,6 @@ class SinkhornOutput(NamedTuple):
       below the convergence threshold.
     inner_iterations: number of iterations that were run between two
       computations of errors.
-
   """
 
   f: Optional[jnp.ndarray] = None
@@ -424,10 +423,6 @@ def transport_cost_at_geom(
       return jnp.sum(self.apply(geom.cost_1.T) * geom.cost_2.T)
     return jnp.sum(self.matrix * other_geom.cost_matrix)
 
-  @property
-  def linear(self) -> bool:  # noqa: D102
-    return isinstance(self.ot_prob, linear_problem.LinearProblem)
-
   @property
   def geom(self) -> geometry.Geometry:  # noqa: D102
     return self.ot_prob.geom
@@ -440,17 +435,10 @@ def a(self) -> jnp.ndarray:  # noqa: D102
   def b(self) -> jnp.ndarray:  # noqa: D102
     return self.ot_prob.b
 
-  @property
-  def linear_output(self) -> bool:  # noqa: D102
-    return True
-
-  # TODO(michalk8): this should be always present
   @property
   def n_iters(self) -> int:  # noqa: D102
     """Returns the total number of iterations that were needed to terminate."""
-    if self.errors is None:
-      return -1
-    return jnp.sum(self.errors > -1) * self.inner_iterations
+    return jnp.sum(self.errors != -1) * self.inner_iterations
 
   @property
   def scalings(self) -> Tuple[jnp.ndarray, jnp.ndarray]:  # noqa: D102

diff --git a/src/ott/solvers/linear/sinkhorn_lr.py b/src/ott/solvers/linear/sinkhorn_lr.py
@@ -175,6 +175,7 @@ class LRSinkhornOutput(NamedTuple):
   errors: jnp.ndarray
   ot_prob: linear_problem.LinearProblem
   epsilon: float
+  inner_iterations: int
   # TODO(michalk8): Optional is an artifact of the current impl., refactor
   reg_ot_cost: Optional[float] = None
 
@@ -205,10 +206,6 @@ def compute_reg_ot_cost(  # noqa: D102
         use_danskin=use_danskin
     )
 
-  @property
-  def linear(self) -> bool:  # noqa: D102
-    return isinstance(self.ot_prob, linear_problem.LinearProblem)
-
   @property
   def geom(self) -> geometry.Geometry:  # noqa: D102
     return self.ot_prob.geom
@@ -222,8 +219,8 @@ def b(self) -> jnp.ndarray:  # noqa: D102
     return self.ot_prob.b
 
   @property
-  def linear_output(self) -> bool:  # noqa: D102
-    return True
+  def n_iters(self) -> int:  # noqa: D102
+    return jnp.sum(self.errors != -1) * self.inner_iterations
 
   @property
   def converged(self) -> bool:  # noqa: D102
@@ -773,6 +770,7 @@ def output_from_state(
         costs=state.costs,
         errors=state.errors,
         epsilon=self.epsilon,
+        inner_iterations=self.inner_iterations,
     )
 
   def _converged(self, state: LRSinkhornState, iteration: int) -> bool:

diff --git a/src/ott/solvers/nn/conjugate_solvers.py b/src/ott/solvers/nn/conjugate_solvers.py
@@ -75,16 +75,18 @@ class FenchelConjugateLBFGS(FenchelConjugateSolver):
     max_iter: maximum number of iterations
     max_linesearch_iter: maximum number of line search iterations
     linesearch_type: type of line search
-    decrease_factor: decrease factor for a backtracking line search
-    ls_method: the line search method
+    linesearch_init: strategy for line search initialization
+    increase_factor: factor by which to increase the step size during
+      the line search
   """
 
   gtol: float = 1e-3
   max_iter: int = 10
   max_linesearch_iter: int = 10
-  linesearch_type: Literal["zoom", "backtracking"] = "backtracking"
-  decrease_factor: float = 0.66
-  ls_method: Literal["wolf", "strong-wolfe"] = "strong-wolfe"
+  linesearch_type: Literal["zoom", "backtracking",
+                           "hager-zhang"] = "backtracking"
+  linesearch_init: Literal["increase", "max", "current"] = "increase"
+  increase_factor: float = 1.5
 
   def solve(  # noqa: D102
       self,
@@ -98,9 +100,9 @@ def solve(  # noqa: D102
         fun=lambda x: f(x) - x.dot(y),
         tol=self.gtol,
         maxiter=self.max_iter,
-        decrease_factor=self.decrease_factor,
         linesearch=self.linesearch_type,
-        condition=self.ls_method,
+        linesearch_init=self.linesearch_init,
+        increase_factor=self.increase_factor,
         implicit_diff=False,
         unroll=False
     )

diff --git a/src/ott/solvers/quadratic/gromov_wasserstein.py b/src/ott/solvers/quadratic/gromov_wasserstein.py
@@ -103,7 +103,7 @@ def primal_cost(self) -> float:
   def n_iters(self) -> int:  # noqa: D102
     if self.errors is None:
       return -1
-    return jnp.sum(self.errors > -1)
+    return jnp.sum(self.errors[:, 0] != -1)
 
 
 class GWState(NamedTuple):