Remove lineax dependency

docs/conf.py

@@ -64,7 +64,6 @@
     "numpy": ("https://numpy.org/doc/stable/", None),
     "jax": ("https://jax.readthedocs.io/en/latest/", None),
     "jaxopt": ("https://jaxopt.github.io/stable", None),
-    "lineax": ("https://docs.kidger.site/lineax/", None),
     "flax": ("https://flax.readthedocs.io/en/latest/", None),
     "optax": ("https://optax.readthedocs.io/en/latest/", None),
     "diffrax": ("https://docs.kidger.site/diffrax/", None),

pyproject.toml

@@ -15,7 +15,6 @@ authors = [
 dependencies = [
     "jax>=0.4.0",
     "jaxopt>=0.8",
-    "lineax>=0.0.7",
     "numpy>=1.20.0",
     "typing_extensions; python_version <= '3.9'",
     "optax>=0.2.4",

src/ott/geometry/regularizers.py

@@ -13,9 +13,7 @@
 # limitations under the License.
 import abc
 import functools
-from typing import Any, Callable, Optional, Tuple, Union
-
-import lineax as lx
+from typing import Any, Callable, Optional, Tuple
 
 import jax
 import jax.numpy as jnp
@@ -185,31 +183,31 @@ class Orthogonal(ProximalOperator):
   def __init__(
       self,
       f: ProximalOperator,
-      A: Optional[Union[jnp.ndarray, lx.AbstractLinearOperator]],
+      A: Optional[jnp.ndarray],
       b: Optional[jnp.ndarray] = None,
       nu: float = 1.0,
   ):
     assert nu > 0.0, nu
     super().__init__()
     self.f = f
     # AA^T = alpha I
-    self.A = lx.MatrixLinearOperator(A) if isinstance(A, jnp.ndarray) else A
+    self.A = A
     self.b = b
     self.nu = nu
 
   def __call__(self, x: jnp.ndarray) -> float:  # noqa: D102
-    z = self.A.mv(x)
+    z = self.A @ x
     if self.b is not None:
       z = z + self.b
     return self.f(z)
 
   def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray:  # noqa: D102
-    w = self.A.mv(v)
+    w = self.A @ v
     if self.b is None:
       tmp = self.f.prox(w, tau * self.nu)
     else:
       tmp = self.f.prox(w + self.b, tau * self.nu) - self.b
-    return v - (1.0 / self.nu) * (self.A.T.mv(w - tmp))
+    return v - (1.0 / self.nu) * (self.A.T @ (w - tmp))
 
   @property
   def is_fully_orthogonal(self) -> bool:
@@ -239,22 +237,22 @@ class Quadratic(ProximalOperator):
       :math:`A`, defined as :math:`A^{\perp} := I - A^T (AA^T)^{-1} A`.
     is_orthogonal: Whether :math:`AA^T = I`.
     is_factor: Whether to factor the matrix :math:`Q` as mentioned above.
-    solver: Linear solver. If :obj:`None`, use :func:`lineax.linear_solve`.
+    solver: Linear solver. If :obj:`None`, use :func:`jnp.linalg.solve`.
   """
 
   def __init__(
       self,
-      A: Optional[Union[jnp.ndarray, lx.AbstractLinearOperator]] = None,
+      A: Optional[jnp.ndarray] = None,
       b: Optional[jnp.ndarray] = None,
       *,
       is_complement: bool = False,
       is_orthogonal: bool = False,
       is_factor: bool = False,
-      solver: Optional[Callable[[lx.AbstractLinearOperator, jnp.ndarray],
+      solver: Optional[Callable[[jnp.ndarray, jnp.ndarray],
                                 jnp.ndarray]] = None,
   ):
     super().__init__()
-    self.A = lx.MatrixLinearOperator(A) if isinstance(A, jnp.ndarray) else A
+    self.A = A
     self.b = b
     self._is_complement = is_complement
     self._is_orthogonal = is_orthogonal
@@ -263,7 +261,7 @@ def __init__(
 
   def __call__(self, x: jnp.ndarray) -> float:  # noqa: D102
     Q = self.Q
-    y = 0.5 * (jnp.dot(x, x) if Q is None else jnp.dot(x, Q.mv(x)))
+    y = 0.5 * (jnp.dot(x, x) if Q is None else jnp.dot(x, Q @ x))
     return y if self.b is None else (y + jnp.dot(x, self.b))
 
   def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray:  # noqa: D102
@@ -274,28 +272,29 @@ def prox(self, v: jnp.ndarray, tau: float = 1.0) -> jnp.ndarray:  # noqa: D102
     if Q is None:
       return (1.0 / (1.0 + tau)) * b
 
-    iden = lx.IdentityLinearOperator(Q.out_structure())
+    n = Q.shape[0]
+    iden = jnp.eye(n)
     if self.is_factor:  # use matrix inversion lemma
       if self.is_complement:
         # eq. 14 in :cite:`klein:24`
         # A_comp = I - A^T(AA^T)^{-1}A
         # prox(v) = (I + tau A_comp^T A_comp)^{-1} (v - tau * b)
         op = iden + tau * (iden - self.A_comp)
-        return (1.0 / (1.0 + tau)) * op.mv(b)
+        return (1.0 / (1.0 + tau)) * (op @ b)
       if self.is_orthogonal:
         # https://en.wikipedia.org/wiki/Woodbury_matrix_identity
         op = iden - (tau / (1.0 + tau)) * (self.A.T @ self.A)
-        return op.mv(b)
+        return op @ b
 
     A = iden + tau * Q
 
     if self._solver is None:
       # use default solver
-      return lx.linear_solve(A, b).value
+      return jnp.linalg.solve(A, b)
     return self._solver(A, b)
 
   @property
-  def A_comp(self) -> Optional[lx.AbstractLinearOperator]:
+  def A_comp(self) -> Optional[jnp.ndarray]:
     r"""Orthogonal complement :math:`A^{\perp}` of :math:`A`."""
     return _complement(
         self.A, self.is_orthogonal
@@ -317,7 +316,7 @@ def is_orthogonal(self) -> bool:
     return self.A is not None and self._is_orthogonal
 
   @property
-  def Q(self) -> Optional[lx.AbstractLinearOperator]:
+  def Q(self) -> Optional[jnp.ndarray]:
     r"""Linear operator :math:`Q`."""
     Q = self.A_comp if self.is_complement else self.A
     if Q is None:
@@ -356,7 +355,7 @@ class SqL2(ProximalOperator):
 
   def __init__(
       self,
-      A: Optional[Union[jnp.ndarray, lx.AbstractLinearOperator]] = None,
+      A: Optional[jnp.ndarray] = None,
       **kwargs: Any,
   ):
     super().__init__()
@@ -498,23 +497,17 @@ def tree_flatten(self):  # noqa: D102
     return (), {"k": self.k}
 
 
-def _invert(A: lx.AbstractLinearOperator) -> lx.MatrixLinearOperator:
-  d = A.out_size()
-  b = jnp.zeros(d)
-
-  solve_fn = jax.vmap(lambda ix: lx.linear_solve(A, b.at[ix].set(1.0)).value)
-  inv = solve_fn(jnp.arange(d))
-  return lx.MatrixLinearOperator(inv)
+def _invert(A: jnp.ndarray) -> jnp.ndarray:
+  return jnp.linalg.inv(A)
 
 
 @functools.partial(jax.jit, static_argnums=1)
-def _complement(
-    A: lx.AbstractLinearOperator, is_orthogonal: bool
-) -> lx.AbstractLinearOperator:
-  iden = lx.IdentityLinearOperator(A.in_structure())
+def _complement(A: jnp.ndarray, is_orthogonal: bool) -> jnp.ndarray:
+  n = A.shape[1]
+  iden = jnp.eye(n)
   if is_orthogonal:
     # AA^T = I
     return iden - (A.T @ A)
 
-  A_inv = _invert(lx.TaggedLinearOperator(A @ A.T, tags={lx.symmetric_tag}))
+  A_inv = _invert(A @ A.T)
   return iden - A.T @ (A_inv @ A)

src/ott/solvers/linear/implicit_differentiation.py

@@ -43,12 +43,12 @@ class ImplicitDiff:
     solver: Callable to compute the solution to a linear problem. The callable
       expects a linear function, a vector, optionally another linear function
       that implements the transpose of that function, and a boolean flag to
-      specify symmetry. This solver is by default one of :class:`lineax.CG` or
-      :class:`lineax.NormalCG` solvers, if the package can be imported, as
-      described in :func:`~ott.solvers.linear.lineax_implicit.solve_lineax`.
+      specify symmetry. This solver is by default a conjugate gradient solver,
+      as described in
+      :func:`~ott.solvers.linear.lineax_implicit.solve_lineax`.
       The :mod:`jax` alternative is described in
       :func:`~ott.solvers.linear.implicit_differentiation.solve_jax_cg`.
-      Note that `lineax` solvers handle better poorly conditioned problems,
+      Note that the CG solver handles better poorly conditioned problems,
       which arise typically when differentiating the solutions of balanced OT
       problems (when ``tau_a==tau_b==1.0``). Relying on
       :func:`~ott.solvers.linear.implicit_differentiation.solve_jax_cg`
@@ -141,7 +141,7 @@ def solve(
     block.
 
     These linear systems are solved using the user-defined ``solver``, using
-    by default :mod:`lineax` solvers when available, or falling back on
+    by default the conjugate gradient solver, or falling back on
     :mod:`jax` when not.
 
     Args:

tests/geometry/regularizers_test.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 from typing import Optional
 
-import lineax as lx
-
 import pytest
 
 import jax
@@ -125,7 +123,7 @@ def test_properties(reg: regularizers.ProximalOperator) -> None:
       assert reg.is_complement == is_complement
       assert reg.is_factor == is_factor
       if reg.is_complement:
-        assert isinstance(reg.A_comp, lx.AbstractLinearOperator)
+        assert isinstance(reg.A_comp, jnp.ndarray)
       else:
         assert reg.A_comp is None
 
@@ -144,7 +142,7 @@ def test_properties(reg: regularizers.ProximalOperator) -> None:
         is_factor=is_factor,
     )
     grad_reg = jax.jit(jax.grad(loss))(reg, x)
-    grad_A = grad_reg.A.as_matrix()
+    grad_A = grad_reg.A
 
     test_properties(reg)
     test_properties(grad_reg)
@@ -200,9 +198,9 @@ def test_orth_use_b(
         is_orthogonal=is_orthogonal,
     )
 
-    iden = lx.IdentityLinearOperator(reg_orth.Q.out_structure())
+    iden = jnp.eye(reg_orth.Q.shape[0])
     A, y = iden + tau * reg_orth.Q, x - tau * b
-    expected = lx.linear_solve(A, y).value
+    expected = jnp.linalg.solve(A, y)
     actual = reg_orth.prox(x, tau)
 
     np.testing.assert_allclose(expected, actual, rtol=1e-3, atol=1e-3)
@@ -238,16 +236,16 @@ def test_l2(self, rng: jax.Array, lam: float, is_complement: bool):
     l2 = regularizers.PostComposition(l2, alpha=lam)
     reg = l2.f.f
 
-    A_ = (reg.A_comp if is_complement else reg.A).as_matrix()
+    A_ = (reg.A_comp if is_complement else reg.A)
     expected_norm1 = 0.5 * lam * jnp.dot(A_ @ x, A_ @ x)
-    expected_norm2 = 0.5 * lam * jnp.dot(x, reg.Q.as_matrix() @ x)
+    expected_norm2 = 0.5 * lam * jnp.dot(x, reg.Q @ x)
 
-    assert reg.A.as_matrix().shape == (k, d)
-    assert reg.Q.as_matrix().shape == (d, d)
+    assert reg.A.shape == (k, d)
+    assert reg.Q.shape == (d, d)
     assert reg.is_factor
     assert reg.is_complement == is_complement
     if is_complement:
-      assert reg.A_comp.as_matrix().shape == (d, d)
+      assert reg.A_comp.shape == (d, d)
 
     np.testing.assert_allclose(expected_norm1, l2(x), rtol=1e-5, atol=1e-5)
     np.testing.assert_allclose(expected_norm2, l2(x), rtol=1e-5, atol=1e-5)
@@ -299,9 +297,9 @@ def test_matrix_inversion_lemma(
         is_factor=True,
         is_complement=is_complement,
     )
-    iden = lx.IdentityLinearOperator(reg.Q.out_structure())
+    iden = jnp.eye(reg.Q.shape[0])
 
-    expected = lx.linear_solve(iden + tau * reg.Q, x - tau * b).value
+    expected = jnp.linalg.solve(iden + tau * reg.Q, x - tau * b)
     actual = reg.prox(x, tau)
 
     np.testing.assert_allclose(expected, actual, rtol=1e-5, atol=1e-5)

tests/problems/linear/potentials_test.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import lineax as lx
-
 import pytest
 
 import jax
@@ -249,7 +247,6 @@ def proj(matrix: jnp.ndarray) -> jnp.ndarray:
       return u.dot(v_h)
 
     def create_cost(A: jnp.ndarray) -> costs.RegTICost:
-      A = lx.MatrixLinearOperator(A)
       orth = regularizers.Orthogonal(reg, A=A)
       return costs.RegTICost(orth, lam=1.0)
 
@@ -286,7 +283,7 @@ def loss(c: costs.RegTICost) -> float:
     expected = (loss_p_delta - loss_m_delta) / (2.0 * eps)
 
     grad_matrix = jax.jit(jax.grad(loss))(cost_fn)
-    grad_matrix = grad_matrix.regularizer.f.A.as_matrix()
+    grad_matrix = grad_matrix.regularizer.f.A
     np.testing.assert_allclose(
         expected, jnp.vdot(delta, grad_matrix), rtol=1e-2, atol=1e-2
     )

tests/tools/soft_sort_test.py

@@ -278,11 +278,7 @@ def test_soft_sort_jacobian(self, rng: jax.Array, implicit: bool):
       pytest.skip(reason="Implicit doesn't work on Python>=3.9 and Linux.")
 
     # Add a ridge when using JAX solvers.
-    try:
-      from ott.solvers.linear import lineax_implicit  # noqa: F401
-      solver_kwargs = {}
-    except ImportError:
-      solver_kwargs = {"ridge_identity": 1e-1, "ridge_kernel": 1e-1}
+    solver_kwargs = {}
 
     b, n = 10, 40
     num_targets = n // 2