diff --git a/docs/dae.md b/docs/dae.md index 4a9202f..db599b9 100644 --- a/docs/dae.md +++ b/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, diff --git a/docs/sdae.md b/docs/sdae.md index fe6a19a..3f0330b 100644 --- a/docs/sdae.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 6cf17fb..91e922f 100644 --- a/pyproject.toml +++ b/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] diff --git a/src/tinydiffeq/dae.py b/src/tinydiffeq/dae.py index ce8d67b..d371b2a 100644 --- a/src/tinydiffeq/dae.py +++ b/src/tinydiffeq/dae.py @@ -77,28 +77,44 @@ 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)) @@ -106,30 +122,7 @@ class LMRootSolver: 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), ) diff --git a/tests/conftest.py b/tests/conftest.py index ac63b78..12e7992 100644 --- a/tests/conftest.py +++ b/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") diff --git a/tests/test_dae.py b/tests/test_dae.py index 3a828f4..716a45a 100644 --- a/tests/test_dae.py +++ b/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(): diff --git a/tests/test_gpu.py b/tests/test_gpu.py index c186766..e8a36f0 100644 --- a/tests/test_gpu.py +++ b/tests/test_gpu.py @@ -526,3 +526,137 @@ def evaluate(rate): for leaf in jax.tree.leaves(result): assert leaf.devices().pop().platform == "gpu" assert bool(jnp.isfinite(leaf)) + + +# The tests above build state with a bare `jnp.asarray(1.0)`, which is float64 +# under the x64 conftest -- so the core ODE/SDE/DAE GPU paths were never +# exercised in single precision. These are float32 end to end. The bar is +# deliberately coarse: no crash, on-device, finite, and right to a few digits. +# Tight float32 agreement is a CPU concern; here the point is that the GPU +# kernels take the same paths and do not produce nonsense. + +F32_DECAY = 0.2 + + +def test_float32_adaptive_ode_runs_on_gpu(): + gpu = gpu_devices()[0] + + def f(x, t, args, p): + return -p * x + + @jax.jit + def run(x_0, p): + # SaveAt defaults to t_1=True, so xs IS the endpoint. + return solve_ode( + f, Tsit5(), 0.0, 1.0, x_0, p=p, dt_0=0.1, controller=IController() + ).xs + + with jax.default_device(gpu): + out = jax.block_until_ready( + run(jnp.asarray(1.0, jnp.float32), jnp.asarray(F32_DECAY, jnp.float32)) + ) + + assert out.dtype == jnp.float32 + assert out.devices().pop().platform == "gpu" + assert bool(jnp.isfinite(out)) + assert abs(float(out) - float(jnp.exp(-F32_DECAY))) < 1e-3 + + +def test_float32_sde_runs_on_gpu(): + gpu = gpu_devices()[0] + + @jax.jit + def run(x_0): + return solve_sde( + lambda x, t, args, p: -F32_DECAY * x, + lambda x, t, args, p: jnp.asarray(0.1, x.dtype) * jnp.ones_like(x), + EulerMaruyama(), + 0.0, + 1.0, + x_0, + key=jax.random.key(0), + n_steps=64, + ).xs + + with jax.default_device(gpu): + out = jax.block_until_ready(run(jnp.asarray(1.0, jnp.float32))) + + assert out.dtype == jnp.float32 + assert out.devices().pop().platform == "gpu" + assert bool(jnp.isfinite(out)) + assert abs(float(out)) < 10.0 # a diverged path, not a tolerance check + + +def test_float32_dae_value_and_grad_run_on_gpu(): + gpu = gpu_devices()[0] + + def f(state, z): + return -F32_DECAY * z + + def g(y, z, t, args, p): + return z - p * y + + def endpoint(p): + return jnp.sum( + solve_semi_explicit_dae( + f, + g, + Tsit5(), + 0.0, + 1.0, + jnp.asarray([1.0], jnp.float32), + jnp.asarray([1.0], jnp.float32), + p=p, + dt_0=0.1, + controller=IController(), + ).ys[-1] + ) + + with jax.default_device(gpu): + value, grad = jax.block_until_ready( + jax.jit(lambda p: (endpoint(p), jax.grad(endpoint)(p)))( + jnp.asarray(1.0, jnp.float32) + ) + ) + + for leaf in (value, grad): + assert leaf.dtype == jnp.float32 + assert leaf.devices().pop().platform == "gpu" + assert bool(jnp.isfinite(leaf)) + + +def test_float32_y_with_float64_z_keeps_the_tight_root_tolerance_on_gpu(): + # dae.py picks root_atol from z_dtype alone: 1e-10 when z has more than 32 + # bits, 1e-6 otherwise. A float32 y with a float64 z therefore gets the + # tight bar, and both dtypes must survive the round trip on device. + gpu = gpu_devices()[0] + + def f(state, z): + return -F32_DECAY * state + + def g(y, z, t, args, p): + return z - jnp.asarray(y, z.dtype) + + @jax.jit + def run(y_0, z_0): + solution = solve_semi_explicit_dae( + f, + g, + Tsit5(), + 0.0, + 1.0, + y_0, + z_0, + dt_0=0.1, + controller=IController(), + ) + return solution.ys[-1], solution.zs[-1] + + with jax.default_device(gpu): + ys, zs = jax.block_until_ready( + run(jnp.asarray([1.0], jnp.float32), jnp.asarray([1.0], jnp.float64)) + ) + + assert ys.dtype == jnp.float32 + assert zs.dtype == jnp.float64 + assert bool(jnp.all(jnp.isfinite(ys))) and bool(jnp.all(jnp.isfinite(zs))) diff --git a/uv.lock b/uv.lock index a526f89..b978581 100644 --- a/uv.lock +++ b/uv.lock @@ -547,14 +547,14 @@ wheels = [ [[package]] name = "nlls-gram" -version = "2.4.0" +version = "2.7.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "jax" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/bc/cd/3195db6a86d1e40b703a12fb796924aa2eeb703ad65eb1ee22585e431edf/nlls_gram-2.4.0.tar.gz", hash = "sha256:d6b7e42b4dc3d05b7a388ccb37b7b09baa5982d10aead72ea010572bf93c4248", size = 1959245, upload-time = "2026-07-21T22:30:04.291Z" } +sdist = { url = "https://files.pythonhosted.org/packages/80/3f/4108c06ceeb8f102f8b6445efc4530aef46902138dd7ee1ffdf7da0f0a5a/nlls_gram-2.7.0.tar.gz", hash = "sha256:1dc3d93792e4d23ff6bd2e13690aaf545e85f76fec23dc4804804b7b5935cec7", size = 230210, upload-time = "2026-07-26T02:13:40.567Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/c3/7d/dd2e404450bdb687669a35ae45ecb15f04e0932bdde547e31722b2d60b53/nlls_gram-2.4.0-py3-none-any.whl", hash = "sha256:d2beceeb15d6e0668d4ad63e4a04c69b7a6a5f88f8863171c77454743930f7c7", size = 76950, upload-time = "2026-07-21T22:30:02.511Z" }, + { url = "https://files.pythonhosted.org/packages/d7/00/4776317cf40e909edd77b79a38e1464b25026d431cbf57699627c2727a6a/nlls_gram-2.7.0-py3-none-any.whl", hash = "sha256:2ff6b83e6028ba266e2f3ca64b001fddf888cf88171e1e1c5a2ded1336e6d531", size = 77552, upload-time = "2026-07-26T02:13:39.213Z" }, ] [[package]] @@ -1227,7 +1227,7 @@ wheels = [ [[package]] name = "tinydiffeq" -version = "2.2.0" +version = "2.3.0" source = { editable = "." } dependencies = [ { name = "jax" }, @@ -1255,7 +1255,7 @@ gpu = [ [package.metadata] requires-dist = [ { name = "jax", specifier = ">=0.7.0" }, - { name = "nlls-gram", specifier = ">=2.4.0" }, + { name = "nlls-gram", specifier = ">=2.7.0" }, ] [package.metadata.requires-dev]