Enable setting of Jacobian tags in ODETerm

diffrax/_adjoint.py

@@ -1,7 +1,7 @@
 import abc
 import functools as ft
 import warnings
-from collections.abc import Callable, Iterable
+from collections.abc import Callable
 from typing import Any, cast
 
 import equinox as eqx
@@ -15,6 +15,7 @@
 from equinox.internal import ω
 
 from ._heuristics import is_sde, is_unsafe_sde
+from ._misc import _frozenset
 from ._saveat import save_y, SaveAt, SubSaveAt
 from ._solver import (
     AbstractItoSolver,
@@ -452,15 +453,6 @@ def _solve(inputs):
     _solve = _solve.__wrapped__  # pyright: ignore[reportFunctionMemberAccess]
 
 
-def _frozenset(x: object | Iterable[object]) -> frozenset[object]:
-    try:
-        iter_x = iter(x)  # pyright: ignore
-    except TypeError:
-        return frozenset([x])
-    else:
-        return frozenset(iter_x)
-
-
 class ImplicitAdjoint(AbstractAdjoint):
     r"""Backpropagate via the [implicit function theorem](https://en.wikipedia.org/wiki/Implicit_function_theorem#Statement_of_the_theorem).
 

diffrax/_misc.py

@@ -1,4 +1,4 @@
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from typing import cast
 
 import jax
@@ -13,6 +13,15 @@
 from ._custom_types import BoolScalarLike, RealScalarLike
 
 
+def _frozenset(x: object | Iterable[object]) -> frozenset[object]:
+    try:
+        iter_x = iter(x)  # pyright: ignore
+    except TypeError:
+        return frozenset([x])
+    else:
+        return frozenset(iter_x)
+
+
 def force_bitcast_convert_type(val, new_type):
     val = jnp.asarray(val)
     result = lax.bitcast_convert_type(val, new_type)

diffrax/_solver/base.py

@@ -14,6 +14,7 @@
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu
+import lineax as lx
 import optimistix as optx
 
 
@@ -214,6 +215,47 @@ class AbstractImplicitSolver(AbstractSolver[_SolverState]):
     root_finder: AbstractVar[optx.AbstractRootFinder]
     root_find_max_steps: AbstractVar[int]
 
+    def _residual_tags(
+        self,
+        jac_tags: frozenset,
+        y_struct,
+        implicit_op_relation: Callable[
+            [lx.AbstractLinearOperator], lx.AbstractLinearOperator
+        ],
+    ) -> frozenset:
+        """Derive residual Jacobian tags from ODE Jacobian tags via lineax composition.
+
+        `implicit_op_relation` is an `op -> op` function that mirrors the solver's
+        `implicit_relation`: given a dummy ODE Jacobian operator `J`, it returns the
+        corresponding residual Jacobian operator (e.g. `I - J` for DIRK, `J - I` for
+        ImplicitEuler). Tags are then read off the composed operator via lineax's
+        operator arithmetic, avoiding any manual tag-mapping logic.
+        """
+        if not jac_tags:
+            return frozenset()
+        dummy_J = lx.FunctionLinearOperator(
+            lambda v: v, y_struct, tags=jac_tags, closure_convert=False
+        )
+        composed = implicit_op_relation(dummy_J)
+        result: set[object] = set()
+        if lx.is_symmetric(composed):
+            result.add(lx.symmetric_tag)
+        if lx.is_positive_semidefinite(composed):
+            result.add(lx.positive_semidefinite_tag)
+        if lx.is_negative_semidefinite(composed):
+            result.add(lx.negative_semidefinite_tag)
+        if lx.is_diagonal(composed):
+            result.add(lx.diagonal_tag)
+        if lx.is_tridiagonal(composed):
+            result.add(lx.tridiagonal_tag)
+        if lx.has_unit_diagonal(composed):
+            result.add(lx.unit_diagonal_tag)
+        if lx.is_lower_triangular(composed):
+            result.add(lx.lower_triangular_tag)
+        if lx.is_upper_triangular(composed):
+            result.add(lx.upper_triangular_tag)
+        return frozenset(result)
+
 
 class AbstractItoSolver(AbstractSolver[_SolverState]):
     """Indicates that when used as an SDE solver that this solver will converge to the

diffrax/_solver/implicit_euler.py

@@ -1,8 +1,10 @@
 from collections.abc import Callable
 from typing import ClassVar, TypeAlias
 
+import jax
 import jax.numpy as jnp
 import jax.tree_util as jtu
+import lineax as lx
 import optimistix as optx
 from equinox.internal import ω
 
@@ -11,7 +13,7 @@
 from .._local_interpolation import LocalLinearInterpolation
 from .._root_finder import with_stepsize_controller_tols
 from .._solution import is_okay, RESULTS
-from .._term import AbstractTerm
+from .._term import AbstractTerm, WrapTerm
 from .base import AbstractAdaptiveSolver, AbstractImplicitSolver
 
 
@@ -24,6 +26,12 @@ def _implicit_relation(z1, nonlinear_solve_args):
     return diff
 
 
+def _implicit_op_relation(J: lx.AbstractLinearOperator) -> lx.AbstractLinearOperator:
+    # Mirrors `_implicit_relation`: residual(z) ≈ h·vf(y0 + z) - z, so the residual
+    # Jacobian is h·J - I. The scalar h is omitted (tag inference is sign-only).
+    return J - lx.IdentityLinearOperator(J.in_structure())
+
+
 class ImplicitEuler(AbstractImplicitSolver, AbstractAdaptiveSolver):
     r"""Implicit Euler method.
 
@@ -82,6 +90,13 @@ def step(
         # (C.f. `AbstractRungeKutta.step`.)
         # If we wanted FSAL then really the correct thing to do would just be to
         # write out a `ButcherTableau` and use `AbstractSDIRK`.
+        _inner = terms.term if isinstance(terms, WrapTerm) else terms
+        y_struct = jax.tree_util.tree_map(
+            lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), y0
+        )
+        residual_tags = self._residual_tags(
+            getattr(_inner, "tags", frozenset()), y_struct, _implicit_op_relation
+        )
         k0 = terms.vf_prod(t0, y0, args, control)
         args = (terms.vf_prod, t1, y0, args, control)
         nonlinear_sol = optx.root_find(
@@ -91,6 +106,7 @@ def step(
             args,
             throw=False,
             max_steps=self.root_find_max_steps,
+            tags=residual_tags,
         )
         k1 = nonlinear_sol.value
         y1 = (y0**ω + k1**ω).ω

diffrax/_solver/runge_kutta.py

@@ -18,6 +18,7 @@
 import jax.lax as lax
 import jax.numpy as jnp
 import jax.tree_util as jtu
+import lineax as lx
 import lineax.internal as lxi
 import numpy as np
 import optimistix as optx
@@ -283,6 +284,14 @@ def _implicit_relation_k(ki, nonlinear_solve_args):
     return diff
 
 
+def _implicit_op_relation_k(
+    J: lx.AbstractLinearOperator,
+) -> lx.AbstractLinearOperator:
+    # Mirrors `_implicit_relation_k`: residual(k) ≈ k - c·vf_prod(y + c·k), so the
+    # residual Jacobian is I - c·J. Scalar c omitted (tag inference is sign-only).
+    return lx.IdentityLinearOperator(J.in_structure()) - J
+
+
 _unused = eqxi.str2jax("unused")  # Sentinel that can be passed into `while_loop` etc.
 
 
@@ -531,6 +540,24 @@ def step(
             tableaus
         )
 
+        # Derive residual Jacobian tags from ODETerm.tags (DIRK: residual = I - c·J).
+        # Only applies to AbstractImplicitSolver subclasses; other RK solvers with
+        # implicit stages (e.g. custom multi-tableau solvers) fall back to frozenset().
+        if implicit_term is not None and isinstance(self, AbstractImplicitSolver):
+            _inner = (
+                implicit_term.term
+                if isinstance(implicit_term, WrapTerm)
+                else implicit_term
+            )
+            y_struct = jtu.tree_map(
+                lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), y0
+            )
+            residual_tags = self._residual_tags(
+                getattr(_inner, "tags", frozenset()), y_struct, _implicit_op_relation_k
+            )
+        else:
+            residual_tags = frozenset()
+
         #
         # We have a choice whether to evaluate `vf` to get vector field evaluations
         # ("values in f-space"), or to evaluate `vf_prod` to get (vector field)-control
@@ -936,7 +963,7 @@ def eval_f_jac():
                         options={},
                         f_struct=jax.eval_shape(lambda: f_pred),
                         aux_struct=None,
-                        tags=frozenset(),
+                        tags=residual_tags,
                     )
 
                 def eval_k_jac():
@@ -947,7 +974,7 @@ def eval_k_jac():
                         options={},
                         f_struct=jax.eval_shape(lambda: k_pred),
                         aux_struct=None,
-                        tags=frozenset(),
+                        tags=residual_tags,
                     )
 
                 if self.calculate_jacobian == CalculateJacobian.every_stage:
@@ -1105,7 +1132,7 @@ def buffers(val):
                     options={},
                     f_struct=jax.eval_shape(lambda: get_implicit(f0)),
                     aux_struct=None,
-                    tags=frozenset(),
+                    tags=residual_tags,
                 )
                 jac_k = _unused
             else:
@@ -1127,7 +1154,7 @@ def buffers(val):
                     options={},
                     f_struct=jax.eval_shape(lambda: y0),
                     aux_struct=None,
-                    tags=frozenset(),
+                    tags=residual_tags,
                 )
         dyn_jac_f, static_jac_f = eqx.partition(jac_f, eqx.is_array)
         dyn_jac_k, static_jac_k = eqx.partition(jac_k, eqx.is_array)

diffrax/_term.py

@@ -24,7 +24,7 @@
     VF,
     Y,
 )
-from ._misc import upcast_or_raise
+from ._misc import _frozenset, upcast_or_raise
 from ._path import AbstractPath
 
 
@@ -189,6 +189,9 @@ class ODETerm(AbstractTerm[_VF, RealScalarLike]):
     """
 
     vector_field: Callable[[RealScalarLike, Y, Args], _VF]
+    tags: frozenset[object] = eqx.field(
+        default_factory=frozenset, converter=_frozenset, static=True
+    )
 
     def vf(self, t: RealScalarLike, y: Y, args: Args) -> _VF:
         out = self.vector_field(t, y, args)
@@ -232,6 +235,11 @@ def _mul(v):
     arguments `(t, y, args)`. `t` is a scalar representing the integration time. `y` is
     the evolving state of the system. `args` are any static arguments as passed to
     [`diffrax.diffeqsolve`][].
+- `tags`: optional lineax tags (e.g. `lineax.symmetric_tag`,
+    `lineax.tridiagonal_tag`) describing structural properties of the vector
+    field Jacobian `∂f/∂y`. Implicit solvers use these to automatically select an
+    appropriate linear solver (e.g. Tridiagonal or Cholesky).
+    Leave empty if unsure.
 """
 
 

test/test_solver.py

@@ -6,6 +6,7 @@
 import jax.numpy as jnp
 import jax.random as jr
 import jax.tree_util as jtu
+import lineax as lx
 import optimistix as optx
 import pytest
 
@@ -502,3 +503,101 @@ def test_adaptive_dt0_milstein(getkey):
     diffrax.diffeqsolve(
         terms, solver, 0, 1, None, 1, stepsize_controller=stepsize_controller
     )
+
+
+# ---------------------------------------------------------------------------
+# Unit tests for AbstractImplicitSolver._residual_tags
+# ---------------------------------------------------------------------------
+
+# Use concrete solvers to exercise the inherited method.
+# negate_J=True  → DIRK (residual ≈ I - c·J)
+# negate_J=False → ImplicitEuler (residual ≈ h·J - I)
+_dirk_solver = diffrax.Kvaerno3()
+_impl_solver = diffrax.ImplicitEuler()
+# y_struct size must be > 1; lineax treats size-1 operators as diagonal.
+_y_struct = jax.ShapeDtypeStruct((5,), jnp.float64)
+
+
+@pytest.mark.parametrize(
+    "input_tags, negate_J, expected_tags",
+    [
+        # nsd → psd for DIRK (I - NSD is PSD), preserved for ImplicitEuler (NSD - I)
+        (
+            frozenset({lx.negative_semidefinite_tag}),
+            True,
+            frozenset({lx.positive_semidefinite_tag, lx.symmetric_tag}),
+        ),
+        (
+            frozenset({lx.negative_semidefinite_tag}),
+            False,
+            frozenset({lx.negative_semidefinite_tag, lx.symmetric_tag}),
+        ),
+        # psd → dropped (sign of residual depends on step size); symmetric preserved
+        (
+            frozenset({lx.positive_semidefinite_tag}),
+            True,
+            frozenset({lx.symmetric_tag}),
+        ),
+        (
+            frozenset({lx.positive_semidefinite_tag}),
+            False,
+            frozenset({lx.symmetric_tag}),
+        ),
+        # unit_diagonal → dropped (I ± J never has unit diagonal in general)
+        (frozenset({lx.unit_diagonal_tag}), True, frozenset()),
+        (frozenset({lx.unit_diagonal_tag}), False, frozenset()),
+        # diagonal → preserved; identity is diagonal+tridiagonal → symmetric too
+        (
+            frozenset({lx.diagonal_tag}),
+            True,
+            frozenset({lx.diagonal_tag, lx.symmetric_tag, lx.tridiagonal_tag}),
+        ),
+        (
+            frozenset({lx.diagonal_tag}),
+            False,
+            frozenset({lx.diagonal_tag, lx.symmetric_tag, lx.tridiagonal_tag}),
+        ),
+        # symmetric → preserved
+        (frozenset({lx.symmetric_tag}), True, frozenset({lx.symmetric_tag})),
+        (frozenset({lx.symmetric_tag}), False, frozenset({lx.symmetric_tag})),
+        # tridiagonal → preserved
+        (frozenset({lx.tridiagonal_tag}), True, frozenset({lx.tridiagonal_tag})),
+        (frozenset({lx.tridiagonal_tag}), False, frozenset({lx.tridiagonal_tag})),
+        # lower_triangular → preserved
+        (
+            frozenset({lx.lower_triangular_tag}),
+            True,
+            frozenset({lx.lower_triangular_tag}),
+        ),
+        (
+            frozenset({lx.lower_triangular_tag}),
+            False,
+            frozenset({lx.lower_triangular_tag}),
+        ),
+        # upper_triangular → preserved
+        (
+            frozenset({lx.upper_triangular_tag}),
+            True,
+            frozenset({lx.upper_triangular_tag}),
+        ),
+        (
+            frozenset({lx.upper_triangular_tag}),
+            False,
+            frozenset({lx.upper_triangular_tag}),
+        ),
+        # empty → empty
+        (frozenset(), True, frozenset()),
+        (frozenset(), False, frozenset()),
+    ],
+)
+def test_residual_tags(input_tags, negate_J, expected_tags):
+    solver = _dirk_solver if negate_J else _impl_solver
+    if negate_J:
+        op_relation = lambda J: lx.IdentityLinearOperator(J.in_structure()) - J
+    else:
+        op_relation = lambda J: J - lx.IdentityLinearOperator(J.in_structure())
+    result = solver._residual_tags(input_tags, _y_struct, op_relation)
+    assert result == expected_tags, (
+        f"input={input_tags!r}, negate_J={negate_J}: "
+        f"got {result!r}, expected {expected_tags!r}"
+    )