tinydiffeq 2.3.0: port to nlls-gram 2.7.0

docs/dae.md

@@ -98,11 +98,7 @@ root_solver = LMRootSolver(
     atol=None,          # 1e-6 float32, 1e-10 float64
     gtol=0.0,           # disabled
     xtol=0.0,           # disabled
-    init_damping=1e-3,
-    linear_solver="auto",
-    jacobian_mode="auto",
-    ad_solver="auto",       # square constraint -> direct
-    ad_solver_penalty=None,
+    solver_options=(),  # nlls-gram constructor defaults
 )
 ```
 
@@ -116,15 +112,25 @@ Root tolerances are independent of the outer controller tolerances. The
 corresponding test. By default, exhausting `root_solver.max_steps` accepts the
 last root iterate and uses its implicit derivative even though nlls retains the
 diagnostic `MAX_STEPS` status. Set `max_steps_is_success=False` to require
-`CONVERGED`; the failed root then has zero implicit tangent. The nlls
-constructor controls are also explicit fields on
-`LMRootSolver`: damping and maximum damping, forward solver and Jacobian mode,
-iterative solver tolerances/preconditioners, AD solver tolerances,
-preconditioner and penalty, solve dtypes, metrics/factories, and recycling.
-Their names and semantics match nlls-gram 2.4. Algebraic roots fix
-`cache_jacobian=False` and `geodesic_acceleration=False`, because each DAE
-stage changes the root problem and the intended path is the ordinary dense LM
-step.
+`CONVERGED`; the failed root then has zero implicit tangent. Everything
+algorithmic runs at nlls-gram's own defaults — a dense `Cholesky()` forward
+solve and `ad_solver=None` for the implicit derivative. For the rare root that
+needs to depart from them, `solver_options` forwards keyword arguments
+verbatim to the `LevenbergMarquardt` constructor:
+
+```python
+from nlls_gram import QR
+
+root_solver = LMRootSolver(solver_options={"linear_solver": QR()})
+```
+
+The names and semantics are nlls-gram's, so they track that package rather
+than being mirrored here. Pass a mapping or key/value pairs; it is normalized
+to a sorted tuple so equal configurations remain hashable and share one
+compiled solver. Algebraic roots fix `cache_jacobian=False` and
+`geodesic_acceleration=False` — each DAE stage changes the root problem and
+the intended path is the ordinary dense LM step — and `solver_options` rejects
+both rather than silently honoring an override.
 
 Every nonlinear root passes `(y, t, p)` through nlls-gram's differentiated parameter
 pytree. Thus it differentiates the defining constraint,

docs/sdae.md

@@ -83,7 +83,7 @@ so JVP/VJP with respect to `y_0` and `p` are pathwise derivatives under common
 random numbers. The key is not differentiable.
 
 `z_0` is a root guess, receives zero tangent, and selects a local root branch.
-Algebraic solves use the same nlls-gram 2.4 `LMRootSolver` configuration and
+Algebraic solves use the same `LMRootSolver` configuration and
 implicit-AD contract as deterministic DAEs; see
 [Nonlinear-solve and AD contract](dae.md#nonlinear-solve-and-ad-contract).
 `MAX_STEPS` is accepted by default; use

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "tinydiffeq"
-version = "2.2.0"
+version = "2.3.0"
 description = "Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX with static shapes and composable AD"
 readme = "README.md"
 license = "MIT"
@@ -29,7 +29,7 @@ classifiers = [
 ]
 dependencies = [
     "jax>=0.7.0",
-    "nlls-gram>=2.4.0",
+    "nlls-gram>=2.7.0",
 ]
 
 [project.urls]

src/tinydiffeq/dae.py

@@ -77,59 +77,52 @@
     Tsit5,
 )
 
+# Fixed for algebraic roots: each DAE stage changes the root problem, so a
+# cached Jacobian is stale and the intended path is the ordinary dense LM step.
+_FIXED_SOLVER_OPTIONS = frozenset({"cache_jacobian", "geodesic_acceleration"})
+
 
 @jax.tree_util.register_dataclass
 @dataclass(frozen=True)
 class LMRootSolver:
     """Configuration for algebraic solves in a semi-explicit DAE.
 
-    The implementation is :class:`nlls_gram.LevenbergMarquardt`. Its
-    shape-adaptive dense default uses ``linear_solver="auto"`` (the normal
-    Cholesky form for a square DAE constraint) and its implicit derivative
-    uses ``ad_solver="auto"``, which resolves every square constraint to the
-    general nonsymmetric direct solve.
-    ``max_steps`` counts nonlinear iterations for one algebraic root and is
-    independent of the integration's time-step ``max_steps``.
-    ``max_steps_is_success=True`` accepts the final iterate when that budget is
-    exhausted; set it to ``False`` to require an nlls ``CONVERGED`` status.
-    ``atol=None`` selects ``1e-6`` in float32 and ``1e-10`` in float64.
-    ``gtol`` and ``xtol`` default to zero (disabled); root tolerances are
-    deliberately independent of the outer integration tolerances.
-
-    The remaining fields pass directly to ``LevenbergMarquardt``. Algebraic
-    residuals do not expose nlls aux, Jacobian caching is disabled because
-    roots change at every DAE stage, and geodesic acceleration is disabled.
+    The implementation is :class:`nlls_gram.LevenbergMarquardt` at its
+    defaults: a dense ``Cholesky()`` forward solve, which takes the normal
+    form for a square DAE constraint, and ``ad_solver=None`` for the implicit
+    derivative, which matches the forward family.
+
+    The fields here are the ones this package owns; all of them reach the nlls
+    ``solve`` rather than its constructor. ``max_steps`` counts nonlinear
+    iterations for one algebraic root and is independent of the integration's
+    time-step ``max_steps``. ``max_steps_is_success=True`` accepts the final
+    iterate when that budget is exhausted; set it to ``False`` to require an
+    nlls ``CONVERGED`` status. ``atol=None`` selects ``1e-6`` in float32 and
+    ``1e-10`` in float64. ``gtol`` and ``xtol`` default to zero (disabled);
+    root tolerances are deliberately independent of the outer integration
+    tolerances.
+
+    ``solver_options`` is the escape hatch for the rare root that needs a
+    non-default algorithm: a mapping (or pairs) forwarded verbatim to the
+    ``LevenbergMarquardt`` constructor, e.g.
+    ``solver_options={"linear_solver": QR()}``. Names and semantics are
+    nlls-gram's, not this package's, so they track it across versions instead
+    of being mirrored here. It is normalized to a sorted tuple so equal
+    configurations stay hashable and share one compiled solver -- an
+    unhashable config silently rebuilds the solver per call, and nlls-gram
+    keys its compiled loop on solver identity, so that would retrace every
+    step. ``cache_jacobian`` and ``geodesic_acceleration`` are fixed to
+    ``False`` and rejected here: each DAE stage changes the root problem, and
+    the intended path is the ordinary dense LM step. Algebraic residuals do
+    not expose nlls aux.
     """
 
     max_steps: int = field(default=8, metadata=dict(static=True))
     max_steps_is_success: bool = field(default=True, metadata=dict(static=True))
     atol: float | None = field(default=None, metadata=dict(static=True))
     gtol: float = field(default=0.0, metadata=dict(static=True))
     xtol: float = field(default=0.0, metadata=dict(static=True))
-    init_damping: float = field(default=1e-3, metadata=dict(static=True))
-    damping_decrease: float = field(default=0.5, metadata=dict(static=True))
-    damping_increase: float = field(default=4.0, metadata=dict(static=True))
-    max_damping: float | None = field(default=None, metadata=dict(static=True))
-    linear_solver: str = field(default="auto", metadata=dict(static=True))
-    jacobian_mode: str = field(default="auto", metadata=dict(static=True))
-    iterative_tol: float = field(default=0.0, metadata=dict(static=True))
-    iterative_atol: float = field(default=0.0, metadata=dict(static=True))
-    iterative_maxiter: int | None = field(default=8, metadata=dict(static=True))
-    dual_preconditioner: Any = field(default=None, metadata=dict(static=True))
-    preconditioner_factory: Any = field(default=None, metadata=dict(static=True))
-    normal_preconditioner: Any = field(default=None, metadata=dict(static=True))
-    whitened_preconditioner: Any = field(default=None, metadata=dict(static=True))
-    ad_solver: str = field(default="auto", metadata=dict(static=True))
-    ad_solver_tol: float | None = field(default=None, metadata=dict(static=True))
-    ad_solver_atol: float = field(default=0.0, metadata=dict(static=True))
-    ad_solver_maxiter: int | None = field(default=None, metadata=dict(static=True))
-    ad_solver_preconditioner: Any = field(default=None, metadata=dict(static=True))
-    ad_solver_penalty: float | None = field(default=None, metadata=dict(static=True))
-    linear_solve_dtype: Any = field(default=None, metadata=dict(static=True))
-    metric_solve_dtype: Any = field(default=None, metadata=dict(static=True))
-    metric: Any = field(default=None, metadata=dict(static=True))
-    metric_factory: Any = field(default=None, metadata=dict(static=True))
-    recycle: Any = field(default=None, metadata=dict(static=True))
+    solver_options: Any = field(default=(), metadata=dict(static=True))
 
     def __post_init__(self):
         if not isinstance(self.max_steps, int) or isinstance(self.max_steps, bool):
@@ -144,12 +137,19 @@ def __post_init__(self):
             raise ValueError("LMRootSolver.gtol must be nonnegative")
         if self.xtol < 0:
             raise ValueError("LMRootSolver.xtol must be nonnegative")
-        if self.init_damping <= 0:
-            raise ValueError("LMRootSolver.init_damping must be positive")
-        if self.damping_decrease <= 0:
-            raise ValueError("LMRootSolver.damping_decrease must be positive")
-        if self.damping_increase <= 0:
-            raise ValueError("LMRootSolver.damping_increase must be positive")
+        try:
+            options = tuple(sorted(dict(self.solver_options).items()))
+        except (TypeError, ValueError) as error:
+            raise TypeError(
+                "LMRootSolver.solver_options must be a mapping or key/value pairs"
+            ) from error
+        fixed = _FIXED_SOLVER_OPTIONS.intersection(name for name, _ in options)
+        if fixed:
+            raise ValueError(
+                "LMRootSolver fixes " + ", ".join(sorted(fixed)) + " for algebraic "
+                "roots; they cannot be set through solver_options"
+            )
+        object.__setattr__(self, "solver_options", options)
 
 
 def _canonicalize_dae_field(fn, name):
@@ -234,32 +234,9 @@ def residual(z, args, root_p):
 
     return LevenbergMarquardt(
         residual,
-        init_damping=config.init_damping,
-        damping_decrease=config.damping_decrease,
-        damping_increase=config.damping_increase,
-        max_damping=config.max_damping,
-        linear_solver=config.linear_solver,
-        jacobian_mode=config.jacobian_mode,
-        iterative_tol=config.iterative_tol,
-        iterative_atol=config.iterative_atol,
-        iterative_maxiter=config.iterative_maxiter,
-        dual_preconditioner=config.dual_preconditioner,
-        preconditioner_factory=config.preconditioner_factory,
-        normal_preconditioner=config.normal_preconditioner,
-        whitened_preconditioner=config.whitened_preconditioner,
-        ad_solver=config.ad_solver,
-        ad_solver_tol=config.ad_solver_tol,
-        ad_solver_atol=config.ad_solver_atol,
-        ad_solver_maxiter=config.ad_solver_maxiter,
-        ad_solver_preconditioner=config.ad_solver_preconditioner,
-        ad_solver_penalty=config.ad_solver_penalty,
-        linear_solve_dtype=config.linear_solve_dtype,
-        metric_solve_dtype=config.metric_solve_dtype,
-        metric=config.metric,
-        metric_factory=config.metric_factory,
         geodesic_acceleration=False,
         cache_jacobian=False,
-        recycle=config.recycle,
+        **dict(config.solver_options),
     )
 
 

tests/conftest.py

@@ -1,3 +1,9 @@
 import jax
 
 jax.config.update("jax_enable_x64", True)
+# XLA:GPU serves float32 dot_general from TF32 tensor cores by default (10-bit
+# mantissa, ~1e-3). The float32 tests compare against closed forms, and the
+# linear-exponential and Markov paths are matmul-heavy, so without this a GPU
+# run disagrees with CPU at ~1e-3 against tolerances set from float32 eps.
+# A no-op on CPU.
+jax.config.update("jax_default_matmul_precision", "highest")

tests/test_dae.py

@@ -1,6 +1,7 @@
 import jax
 import jax.numpy as jnp
 import pytest
+from nlls_gram import QR, Cholesky
 
 from tinydiffeq import (
     RK4,
@@ -39,53 +40,60 @@ def solve_linear(p, y_0, z_0, solver, save_at, **kwargs):
     )
 
 
-def test_lm_root_solver_uses_nlls_24_defaults_and_forwards_options():
+def test_lm_root_solver_uses_nlls_defaults_and_forwards_options():
     def constraint(y, z, t, args, p):
         return z - y
 
     defaults = _build_algebraic_solver(constraint, LMRootSolver(), False)
     assert LMRootSolver().max_steps_is_success
-    assert defaults.linear_solver == "auto"
+    # Everything algorithmic is nlls-gram's default; only the two invariants
+    # this package owns are pinned.
+    assert isinstance(defaults.linear_solver, Cholesky)
     assert defaults.jacobian_mode == "auto"
-    assert defaults.ad_solver == "auto"
-    assert defaults.ad_solver_penalty is None
+    assert defaults.ad_solver is None
     assert not defaults.cache_jacobian
     assert not defaults.geodesic_acceleration
 
     configured = _build_algebraic_solver(
         constraint,
         LMRootSolver(
-            init_damping=2e-3,
-            damping_decrease=0.4,
-            damping_increase=3.0,
-            max_damping=10.0,
-            linear_solver="qr",
-            jacobian_mode="rev",
-            iterative_tol=1e-4,
-            iterative_atol=1e-6,
-            iterative_maxiter=17,
-            ad_solver="augmented_qr",
-            ad_solver_tol=1e-5,
-            ad_solver_atol=1e-7,
-            ad_solver_maxiter=13,
-            ad_solver_penalty=1e-8,
+            solver_options={
+                "init_damping": 2e-3,
+                "damping_decrease": 0.4,
+                "damping_increase": 3.0,
+                "jacobian_mode": "rev",
+                "linear_solver": QR(),
+            }
         ),
         False,
     )
     assert configured.init_damping == 2e-3
     assert configured.damping_decrease == 0.4
     assert configured.damping_increase == 3.0
-    assert configured.max_damping == 10.0
-    assert configured.linear_solver == "qr"
     assert configured.jacobian_mode == "rev"
-    assert configured.iterative_tol == 1e-4
-    assert configured.iterative_atol == 1e-6
-    assert configured.iterative_maxiter == 17
-    assert configured.ad_solver == "augmented_qr"
-    assert configured.ad_solver_tol == 1e-5
-    assert configured.ad_solver_atol == 1e-7
-    assert configured.ad_solver_maxiter == 13
-    assert configured.ad_solver_penalty == 1e-8
+    assert isinstance(configured.linear_solver, QR)
+    # The invariants survive a populated pass-through.
+    assert not configured.cache_jacobian
+    assert not configured.geodesic_acceleration
+
+
+def test_lm_root_solver_options_normalize_and_reject_fixed_keys():
+    # A mapping and the equivalent pairs must compare and hash equal, so
+    # _cached_algebraic_solver shares one compiled solver between them.
+    mapping = LMRootSolver(solver_options={"jacobian_mode": "rev", "init_damping": 0.1})
+    pairs = LMRootSolver(
+        solver_options=(("init_damping", 0.1), ("jacobian_mode", "rev"))
+    )
+    assert mapping == pairs
+    assert hash(mapping) == hash(pairs)
+    # Re-wrapping an already-normalized value is idempotent.
+    assert LMRootSolver(solver_options=mapping.solver_options) == mapping
+
+    for fixed in ("cache_jacobian", "geodesic_acceleration"):
+        with pytest.raises(ValueError, match=fixed):
+            LMRootSolver(solver_options={fixed: True})
+    with pytest.raises(TypeError, match="mapping or key/value pairs"):
+        LMRootSolver(solver_options=5)
 
 
 def test_max_steps_policy_and_strict_batched_derivative():