From 53ba5e5da15ae7392ec8c0cb04051d51aa540b82 Mon Sep 17 00:00:00 2001 From: marcocuturi Date: Fri, 22 May 2026 13:59:01 +0200 Subject: [PATCH 1/2] Remove lineax dependency, use JAX-native linear solvers Replace all lineax usage with plain JAX operations: - lineax_implicit.py: JAX-native CG via jax.lax.while_loop - regularizers.py: plain ndarray operations instead of lineax operators - Remove lineax from pyproject.toml dependencies - Update tests and docs accordingly lineax's equinox closure conversion is incompatible with JAX >= 0.10 when called inside VJP backward passes, causing 39 test failures. The JAX-native implementation avoids this entirely. Co-Authored-By: Claude Opus 4.6 (1M context) --- docs/conf.py | 1 - pyproject.toml | 1 - src/ott/geometry/regularizers.py | 57 ++++++------ .../linear/implicit_differentiation.py | 10 +- src/ott/solvers/linear/lineax_implicit.py | 91 ++++++++++--------- tests/geometry/regularizers_test.py | 24 +++-- tests/problems/linear/potentials_test.py | 5 +- tests/tools/soft_sort_test.py | 6 +- 8 files changed, 91 insertions(+), 104 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index 2e89cc882..f56dc775c 100644 --- a/docs/conf.py +++ b/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), diff --git a/pyproject.toml b/pyproject.toml index a4bff1c53..0e31b5771 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/ott/geometry/regularizers.py b/src/ott/geometry/regularizers.py index 9d6cb56b3..50f087cc9 100644 --- a/src/ott/geometry/regularizers.py +++ b/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,7 +183,7 @@ 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, ): @@ -193,23 +191,23 @@ def __init__( 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,19 @@ 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()) + 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) diff --git a/src/ott/solvers/linear/implicit_differentiation.py b/src/ott/solvers/linear/implicit_differentiation.py index ef3633714..0f64b4adc 100644 --- a/src/ott/solvers/linear/implicit_differentiation.py +++ b/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: diff --git a/src/ott/solvers/linear/lineax_implicit.py b/src/ott/solvers/linear/lineax_implicit.py index fccb502ab..d3ccc17b1 100644 --- a/src/ott/solvers/linear/lineax_implicit.py +++ b/src/ott/solvers/linear/lineax_implicit.py @@ -11,38 +11,53 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Callable, Optional, TypeVar - -import equinox as eqx -import lineax as lx -from jaxtyping import Array, Float, PyTree +from typing import Any, Callable, Optional import jax import jax.numpy as jnp -import jax.tree_util as jtu -_T = TypeVar("_T") -_FlatPyTree = tuple[list[_T], jtu.PyTreeDef] +__all__ = ["solve_lineax"] -__all__ = ["CustomTransposeLinearOperator", "solve_lineax"] +def _cg( + matvec: Callable[[jnp.ndarray], jnp.ndarray], + b: jnp.ndarray, + *, + rtol: float = 1e-6, + atol: float = 1e-6, + maxiter: Optional[int] = None, +) -> jnp.ndarray: + """Conjugate gradient solver using jax.lax.while_loop.""" + if maxiter is None: + maxiter = 10 * b.shape[0] -class CustomTransposeLinearOperator(lx.FunctionLinearOperator): - """Implement a linear operator that can specify its transpose directly.""" - fn: Callable[[PyTree[Float[Array, "..."]]], PyTree[Float[Array, "..."]]] - fn_t: Callable[[PyTree[Float[Array, "..."]]], PyTree[Float[Array, "..."]]] - input_structure: _FlatPyTree[jax.ShapeDtypeStruct] = eqx.field(static=True) - input_structure_t: _FlatPyTree[jax.ShapeDtypeStruct] = eqx.field(static=True) - tags: frozenset[object] + b_norm = jnp.linalg.norm(b) + tol = jnp.maximum(atol, rtol * b_norm) - def __init__(self, fn, fn_t, input_structure, input_structure_t, tags=()): - super().__init__(fn, input_structure, tags) - self.fn_t = eqx.filter_closure_convert(fn_t, input_structure_t) - self.input_structure_t = input_structure_t + x0 = jnp.zeros_like(b) + r0 = b + p0 = r0 + rtr0 = jnp.vdot(r0, r0) - def transpose(self): - """Provide custom transposition operator from function.""" - return lx.FunctionLinearOperator(self.fn_t, self.input_structure_t) + def cond_fun(state): + _, _, _, rtr, k = state + return (jnp.sqrt(rtr) > tol) & (k < maxiter) + + def body_fun(state): + x, r, p, rtr, k = state + Ap = matvec(p) + alpha = rtr / jnp.vdot(p, Ap) + x_new = x + alpha * p + r_new = r - alpha * Ap + rtr_new = jnp.vdot(r_new, r_new) + beta = rtr_new / rtr + p_new = r_new + beta * p + return x_new, r_new, p_new, rtr_new, k + 1 + + x, _, _, _, _ = jax.lax.while_loop( + cond_fun, body_fun, (x0, r0, p0, rtr0, 0) + ) + return x def solve_lineax( @@ -50,29 +65,25 @@ def solve_lineax( b: jnp.ndarray, lin_t: Optional[Callable] = None, symmetric: bool = False, - nonsym_solver: Optional[lx.AbstractLinearSolver] = None, + nonsym_solver: Optional[Any] = None, ridge_identity: float = 0.0, ridge_kernel: float = 0.0, **kwargs: Any ) -> jnp.ndarray: - """Wrapper around lineax solvers. + """Solve a linear system using conjugate gradients. Args: lin: Linear operator b: vector. Returned `x` is such that `lin(x)=b` lin_t: Linear operator, corresponding to transpose of `lin`. symmetric: whether `lin` is symmetric. - nonsym_solver: solver used when handling non-symmetric cases. Note that - :class:`~lineax.Normal` with :class:`~lineax.CG` - is used by default in the symmetric case. + nonsym_solver: unused, kept for API compatibility. ridge_kernel: promotes zero-sum solutions. Only use if `tau_a = tau_b = 1.0` ridge_identity: handles rank deficient transport matrices (this happens typically when rows/cols in cost/kernel matrices are collinear, or, equivalently when two points from either measure are close). - kwargs: arguments passed to :class:`~lineax.AbstractLinearSolver` linear - solver. + kwargs: arguments passed to the CG solver (rtol, atol, maxiter). """ - input_structure = jax.eval_shape(lambda: b) kwargs.setdefault("rtol", 1e-6) kwargs.setdefault("atol", 1e-6) @@ -85,16 +96,8 @@ def solve_lineax( lin_reg, lin_t_reg = lin, lin_t if symmetric: - solver = lx.CG(**kwargs) - fn_operator = lx.FunctionLinearOperator( - lin_reg, input_structure, tags=lx.positive_semidefinite_tag - ) - return lx.linear_solve(fn_operator, b, solver).value - # In the non-symmetric case, use NormalCG by default, but consider - # user defined choice of alternative lx solver. - if nonsym_solver is None: - nonsym_solver = lx.Normal(lx.CG(**kwargs)) - fn_operator = CustomTransposeLinearOperator( - lin_reg, lin_t_reg, input_structure, input_structure - ) - return lx.linear_solve(fn_operator, b, nonsym_solver).value + return _cg(lin_reg, b, **kwargs) + # Non-symmetric: solve normal equations A^T A x = A^T b + normal_matvec = lambda x: lin_t_reg(lin_reg(x)) + normal_b = lin_t_reg(b) + return _cg(normal_matvec, normal_b, **kwargs) diff --git a/tests/geometry/regularizers_test.py b/tests/geometry/regularizers_test.py index 55eccfb63..accb6798f 100644 --- a/tests/geometry/regularizers_test.py +++ b/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) diff --git a/tests/problems/linear/potentials_test.py b/tests/problems/linear/potentials_test.py index 778949251..225129419 100644 --- a/tests/problems/linear/potentials_test.py +++ b/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 ) diff --git a/tests/tools/soft_sort_test.py b/tests/tools/soft_sort_test.py index a331d5da7..ca30f80c1 100644 --- a/tests/tools/soft_sort_test.py +++ b/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 From aa532e29ea8f38b5a4f6852bcaf78d1883475059 Mon Sep 17 00:00:00 2001 From: marcocuturi Date: Fri, 22 May 2026 23:47:57 +0200 Subject: [PATCH 2/2] Fix yapf formatting in regularizers.py Co-Authored-By: Claude Opus 4.6 (1M context) --- src/ott/geometry/regularizers.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/ott/geometry/regularizers.py b/src/ott/geometry/regularizers.py index 50f087cc9..b06aa1952 100644 --- a/src/ott/geometry/regularizers.py +++ b/src/ott/geometry/regularizers.py @@ -502,9 +502,7 @@ def _invert(A: jnp.ndarray) -> jnp.ndarray: @functools.partial(jax.jit, static_argnums=1) -def _complement( - A: jnp.ndarray, is_orthogonal: bool -) -> jnp.ndarray: +def _complement(A: jnp.ndarray, is_orthogonal: bool) -> jnp.ndarray: n = A.shape[1] iden = jnp.eye(n) if is_orthogonal: