diff --git a/README.md b/README.md index f8e5296..c716e9a 100644 --- a/README.md +++ b/README.md @@ -8,71 +8,27 @@ [![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX: fixed-step Euler/RK4, -adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs, -and Euler–Maruyama for Itô SDEs and semi-explicit index-1 SDAEs. The package -also includes primal, vmap-friendly finite-state DTMC and CTMC simulators with -sequential and associative parallel-prefix execution. Deterministic probability -forecasts are differentiable in the initial mass and include DTMC matrix powers, -dense CTMC exponentials, and matrix-free Arnoldi/Krylov actions over probability -pytrees. The same dense and matrix-free backends are available directly through -`solve_linear_ode` for any fixed homogeneous linear array or pytree operator; -`jvp_linear_ode` and `vjp_linear_ode` apply the exact initial-state tangent and -adjoint exponential actions without differentiating Arnoldi orthogonalization. -Fixed stepping and the default adaptive path use bounded `lax.scan` loops with -exactly `max_steps` attempt slots. Shapes stay static as tolerances or curvature -change, and these solves support forward mode, reverse mode, and -reverse-over-forward. Adaptive ODE and DAE solves may instead select -`adaptive_loop="forward"`: a dynamic `lax.while_loop` that executes only actual -attempts and supports JVP and nested forward AD, but not reverse mode. A vmapped -forward loop runs until its slowest lane finishes. +adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 DAEs, and +fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and semi-explicit +index-1 SDAEs. Solves run in bounded `lax.scan` loops with static shapes and +compose with `jit`, `vmap`, forward mode, reverse mode, and +reverse-over-forward. Finite-state Markov simulation, probability forecasts, +and general fixed homogeneous linear solves (dense or matrix-free Krylov +exponential actions, after SciML's +[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/)) +round out the package. This is a deliberately small, jvp/vjp-friendly package. Rodas5P is a JAX -adaptation of Steinebach's method and follows SciML's +adaptation of Steinebach's method following SciML's [`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock) -implementation. Use [diffrax](https://docs.kidger.site/diffrax/) or +implementation, and DAE algebraic roots delegate both the primal solve and +the implicit derivative to +[`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/). Use +[diffrax](https://docs.kidger.site/diffrax/) or [SciML](https://docs.sciml.ai/DiffEqDocs/stable/) if you need general mass -matrices, fully implicit or higher-index DAEs, events, continuous solution -objects, sparse/Krylov linear solvers for ODE/DAE stages, or specialized -adjoints. Initial DAE consistency and explicit DAE stages use `nlls-gram`; -the same nlls solve supplies the square root's implicit derivative, whose -default is a direct nonsymmetric `LU()` solve. `LMRootSolver` requires -residual-only stopping (`gtol=xtol=0`) and accepts only `CONVERGED` roots whose -Euclidean residual norm is below the root `atol`. Its -`max_steps_is_success` field remains for source compatibility but does not make -`MAX_STEPS` a valid DAE root. - -The linear exponential-action API follows SciML -[`ExponentialUtilities.expv`](https://docs.sciml.ai/ExponentialUtilities/stable/expv/). -It includes fixed and residual-controlled adaptive matrix-free time slicing; -the latter keeps the Krylov dimension static for predictable JAX compilation. -SciML's -[`ExponentialIntegrators.jl`](https://docs.sciml.ai/ExponentialIntegrators/stable/) -is the reference for the broader nonlinear exponential-integrator family. - -## 2.4.0 migration note - -- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit - fixed-step ODEs. Every query must align with a knot; adaptive methods, - Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode. -- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts, - including rejections. DAE results additionally expose `num_root_solves` and - `num_root_steps`; `num_accepted` retains its existing meaning. -- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an - actual-work loop. It supports primal, JVP, and nested forward AD but not - reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable - default. Under `vmap`, the forward loop runs to the slowest lane. -- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for - locally unique algebraic branches; `predictor="previous"` remains the - default. - -DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root -solve and implicit derivative; square implicit AD defaults to direct `LU()`. -Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so -`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source -compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid -root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they -relied on budget exhaustion, increase the root budget or adjust the residual -tolerance instead. +matrices, fully implicit or higher-index DAEs, adaptive SDE stepping, events, +continuous solution objects, sparse/Krylov ODE/DAE stages, or specialized +adjoints. ## Install @@ -91,12 +47,8 @@ uv add tinydiffeq "jax[cuda13]" The vector field may take `(x)`, `(x, t)`, `(x, t, args)`, or `(x, t, args, p)` — always in that order. `args` is pass-through data (not an -AD target by convention); `p` holds differentiable parameters (any pytree). -The state may also be any JAX pytree. It must contain at least one leaf, and -every leaf must be a nonempty real floating array with the same dtype; vector -fields and `project` preserve that structure. Output keeps the structure and -adds the saved-time axis to each -leaf. +AD target by convention); `p` holds differentiable parameters, and the state +may be any pytree of same-dtype real floating arrays. ```python import jax @@ -122,29 +74,61 @@ print(sol.xs) # states on the grid print(sol.ok) # reached t_1 with every requested output valid? ``` -`IController()` and `PIController()` choose tolerances from `x_0.dtype`: -`rtol=1e-4, atol=1e-6` for float32 and `rtol=1e-7, atol=1e-9` for -float64. Pass explicit values when tolerances are part of your model's -scientific specification. The default `dt_min` is -`10 * finfo(dtype).eps * max(1, abs(t_1))`. - -`max_steps` is the total internal **attempt budget**: accepted steps plus -rejections. It is not normally the number of returned times. Endpoint mode -returns one time/state, `SaveAt(ts=...)` returns the requested grid, and -`SaveAt(steps=True)` returns the initial state and accepted internal steps as -a contiguous prefix of `max_steps + 1` rows. The remaining rows repeat the -last accepted state by default; `sol.accepted` distinguishes data from -padding. Rejected attempts never appear in the returned trajectory. -`sol.num_steps` reports the number of attempts actually made, while -`sol.num_accepted` excludes rejections. - -`SaveAt(ts=...)` also accepts a Python sequence. These are observation times: -the adaptive controller still chooses its own internal mesh. Explicit methods -use cubic Hermite interpolation; Rodas5P uses its published stiff-aware -fourth-order continuous extension. For an explicit fixed-step ODE, -`SaveAt(ts=..., exact=True)` instead requires every requested time to be an -internal knot and gathers the stored state directly. Exact mode does not apply -to adaptive ODEs, Rodas5P, DAEs, SDEs, or SDAEs. +`max_steps` is the internal attempt budget (accepted plus rejected steps), +not the number of returned times: `SaveAt` picks the endpoint, a fixed +interpolation grid, or the padded accepted-step prefix, so output shapes +never depend on how many steps the controller took. Omitted controller +tolerances follow the state dtype (`1e-4`/`1e-6` in float32, +`1e-7`/`1e-9` in float64). + +## SDEs with first-class noise + +`solve_sde` integrates diagonal-noise Itô SDEs with `EulerMaruyama` (strong +order 0.5), `Milstein` (1.0, commutative diagonal noise), or `SRA1` (1.5, +additive noise). An Ornstein–Uhlenbeck process under SRA1: + +```python +from tinydiffeq import solve_sde, SRA1 + +theta, sigma, n = 1.0, 0.5, 256 + + +def ou_drift(x): + return -theta * x + + +def ou_diffusion(x): + return sigma * jnp.ones_like(x) + + +sol = solve_sde( + ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, jnp.asarray(1.0), + key=jax.random.key(0), n_steps=n, +) +``` + +The noise realization can also be passed explicitly — the same pytree +`sample_noise` would draw, now inspectable, storable data that is +differentiable like any other input: + +```python +x_0 = jnp.asarray(1.0) +noise = SRA1().sample_noise(x_0, jax.random.key(0), n, jnp.asarray(1.0 / n), x_0.dtype) +same_sol = solve_sde( + ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n +) # bit-identical to the key= call +d_endpoint_d_noise = jax.grad( + lambda noise: solve_sde( + ou_drift, ou_diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n + ).xs +)(noise) +``` + +A fixed key (or fixed noise) pins the whole path, so gradients with respect +to `x_0`, `p`, and `noise` are pathwise derivatives under common random +numbers — the setup simulation-based estimators want. `vmap` over +trajectories with per-trajectory `x_0` and noise composes with `jit` and +`grad`. ## Semi-explicit DAEs @@ -152,7 +136,7 @@ For a square index-1 system `dy/dt = f(y, z, t, args, p)` and `0 = g(y, z, t, args, p)`: ```python -from tinydiffeq import IController, Rodas5P, Tsit5, solve_semi_explicit_dae +from tinydiffeq import solve_semi_explicit_dae def dae_f(y, z, t, args, p): @@ -171,36 +155,17 @@ dae_sol = solve_semi_explicit_dae( controller=IController(), max_steps=128, ) print(dae_sol.ys, dae_sol.zs, dae_sol.aux["flow"]) - -# One initial nonlinear consistency solve, then linear Rodas5P stages. -stiff_dae_sol = solve_semi_explicit_dae( - dae_f, dae_g, Rodas5P(), 0.0, 1.0, - jnp.asarray(1.0), jnp.asarray(0.5), - p=jnp.asarray(2.0), dt_0=0.1, - controller=IController(), max_steps=128, -) ``` `z_0` is a guess and is made consistent automatically. RK4 and Tsit5 restore -the algebraic root at every stage. Rodas5P performs no nonlinear solves after -initialization: it advances the corresponding block mass-matrix system using -one reused LU factorization per attempt. Differential fields may return a -floating saved-aux pytree stored at accepted nodes and interpolated on requested -deterministic grids. Algebraic equations may separately return internal context -passed to the dynamics. On the default bounded path, JVP, VJP, and -reverse-over-forward propagate through both implicit initialization and the -time integrator. See the -[DAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) -for root controls, `SaveAt`, and scope limits. - -DAE solutions expose `num_steps`, `num_root_solves`, and `num_root_steps` as -logical per-trajectory work counters. Explicit methods default to reusing the -previous algebraic root as the next stage guess; `LMRootSolver(predictor="secant")` -is an opt-in continuation predictor for locally unique root branches. - -Fixed-step semi-explicit Itô SDAEs use the corresponding -`solve_semi_explicit_sdae` interface with `EulerMaruyama`, a PRNG key, and -`n_steps`; see the [SDAE documentation](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/). +the algebraic root at every stage through `nlls-gram`, which also supplies +the root's implicit derivative; `Rodas5P()` instead performs one initial +consistency solve and then advances the block mass-matrix system with one +reused LU factorization per attempt — the stiff path. Stochastic +semi-explicit systems use `solve_semi_explicit_sdae` with `EulerMaruyama` or +`SRA1`. See the +[DAE](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) and +[SDAE](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) docs. ## Gradients through the solve @@ -212,21 +177,16 @@ def endpoint(p): max_steps=512, ).xs -jax.grad(endpoint)(jnp.asarray(1.3)) # reverse mode +jax.grad(endpoint)(jnp.asarray(1.3)) # reverse mode jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),)) # forward mode -jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])( - jnp.asarray(1.3) -) # reverse-over-forward ``` The step-size controller is wrapped in `stop_gradient` (accept/reject is -non-differentiable either way, and the error-ratio power blows up at exactly -zero error); states differentiate through the solver stages on the realized, -frozen mesh. In particular, adaptive `SaveAt(steps=True)` does not include mesh -motion in its time or state derivatives. See the +non-differentiable either way); states differentiate through the solver +stages on the realized, frozen mesh. See the [docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design -contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key -semantics, and the package API. +contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE +noise semantics, and the package API. ## License diff --git a/benchmarks/README.md b/benchmarks/README.md index 670ecd4..1f887d8 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -1,5 +1,10 @@ # Benchmarks +`gpu_trajectories.py` measures vmapped trajectory ensembles (SRA1/EM SDEs and +fixed RK4, primal and reverse-mode, explicit or key-drawn noise) and is meant +for GPU nodes; run `python -m benchmarks.gpu_trajectories --help` for the +sweep axes. Committed results live under `results/`. + Run the opt-in CPU suite with: ```bash diff --git a/benchmarks/export_growth_policies.py b/benchmarks/export_growth_policies.py new file mode 100644 index 0000000..c894975 --- /dev/null +++ b/benchmarks/export_growth_policies.py @@ -0,0 +1,126 @@ +"""Train the kernels growth policies and export their weights for benchmarks. + +Runs the kernels package's neoclassical and stochastic growth solvers +in-process (the package keeps trained nnx modules in memory only), extracts +the MLP weights and structural parameters as plain arrays, verifies a pure-jnp +reimplementation of the policy forward pass against the nnx module, and writes +`benchmarks/policies/.npz` consumed by `gpu_trajectories.py`. + +Run from the kernels venv (it owns flax and the model modules): + + uv run --project /Users/jlperla/GitHub/kernels python \ + benchmarks/export_growth_policies.py +""" + +import dataclasses +import json +import sys +from pathlib import Path + +sys.path.insert(0, "/Users/jlperla/GitHub/kernels") + +import jax +import jax.numpy as jnp +import numpy as np + +OUTPUT_DIR = Path(__file__).parent / "policies" + + +def unit_softplus(raw): + log_two = jnp.log(jnp.asarray(2.0, dtype=raw.dtype)) + return jax.nn.softplus(2.0 * log_two * raw) / log_two + + +def extract(policy, structural): + weights = {} + for index, layer in enumerate(policy.layers): + weights[f"layer{index}_kernel"] = np.asarray(layer.kernel.value) + weights[f"layer{index}_bias"] = np.asarray(layer.bias.value) + weights["head_kernel"] = np.asarray(policy.head.kernel.value) + weights["head_bias"] = np.asarray(policy.head.bias.value) + metadata = { + "num_layers": len(policy.layers), + "homothetic": bool(policy.homothetic), + "log_z_scale": float(getattr(policy, "log_z_scale", 0.0)), + "log_k_scale": float(getattr(policy, "log_k_scale", 0.0)), + "k_scale": float(getattr(policy, "k_scale", 1.0)), + "structural": { + key: ( + np.asarray(value).item() + if np.asarray(value).ndim == 0 + else np.asarray(value).tolist() + ) + for key, value in dataclasses.asdict(structural).items() + }, + } + return weights, metadata + + +def pure_policy(weights, metadata, gelu_approximate): + num_layers = metadata["num_layers"] + delta = metadata["structural"]["delta"] + + def forward(log_state): + log_z, log_k = log_state + if metadata["homothetic"]: + h = jnp.atleast_1d(log_k - log_z) + anchor = delta * jnp.exp(log_k) + else: + h = jnp.stack( + [log_z - metadata["log_z_scale"], log_k - metadata["log_k_scale"]] + ) + anchor = delta * metadata["k_scale"] + for index in range(num_layers): + h = jax.nn.gelu( + h @ weights[f"layer{index}_kernel"] + weights[f"layer{index}_bias"], + approximate=gelu_approximate, + ) + raw = h @ weights["head_kernel"] + weights["head_bias"] + return anchor * unit_softplus(raw)[0] + + return forward + + +def verify_and_save(name, policy, structural): + weights, metadata = extract(policy, structural) + states = jnp.stack( + [ + 0.3 * jax.random.normal(jax.random.key(0), (256,)), + 0.3 * jax.random.normal(jax.random.key(1), (256,)), + ], + axis=1, + ) + reference = jax.vmap(lambda s: policy(s, structural))(states) + gelu_approximate = None + for candidate in (True, False): + ours = jax.vmap(pure_policy(weights, metadata, candidate))(states) + gap = float(jnp.max(jnp.abs(ours - reference))) + if gap < 1e-10: + gelu_approximate = candidate + break + if gelu_approximate is None: + raise AssertionError(f"{name}: pure-jnp policy mismatch, max gap {gap}") + metadata["gelu_approximate"] = gelu_approximate + OUTPUT_DIR.mkdir(exist_ok=True) + np.savez(OUTPUT_DIR / f"{name}.npz", **weights) + (OUTPUT_DIR / f"{name}.json").write_text(json.dumps(metadata, indent=2) + "\n") + print( + f"{name}: exported {sum(w.size for w in weights.values())} params, " + f"gelu_approximate={gelu_approximate}, parity gap < 1e-10" + ) + + +def main(): + from neoclassical_growth_nn_recursive import neoclassical_growth_nn_recursive + + results, extras = neoclassical_growth_nn_recursive(use_float64=True) + verify_and_save("neoclassical_growth", extras["model"], extras["parameters"]) + + from stochastic_growth_nn_recursive import stochastic_growth_nn_recursive + + results, extras = stochastic_growth_nn_recursive(use_float64=True) + verify_and_save("stochastic_growth", extras["model"], extras["parameters"]) + + +if __name__ == "__main__": + main() diff --git a/benchmarks/gpu_trajectories.py b/benchmarks/gpu_trajectories.py new file mode 100644 index 0000000..852f951 --- /dev/null +++ b/benchmarks/gpu_trajectories.py @@ -0,0 +1,590 @@ +"""Trajectory-ensemble benchmarks for the kernels use case. + +Measures vmapped, jitted solves over B trajectories with per-trajectory x_0 — +SRA1/EulerMaruyama SDEs with explicit or key-drawn noise, and fixed-step RK4 +ODEs — in primal mode and as jit(grad) of a scalar residual with respect to +(x_0, p, noise). n_steps defaults include 31 and 127 so both a scan XLA may +unroll and one it will not are covered. Beyond the synthetic OU/MLP drifts, +`--drifts neoclassical stochastic-growth` load trained kernels-package +investment policies (see export_growth_policies.py) so the measured workload +is the real recursive-NN growth model, with the policy weights as the +differentiable parameters. + +`--remat` sweeps the reverse-mode memory/compute trade (the axis where +diffrax's default RecursiveCheckpointAdjoint differs from a plain scan): +`none` is JAX's native O(n)-memory scan rule, `solve` wraps the whole +per-trajectory solve in jax.checkpoint, and `chunked` is a script-local +sqrt(n)-nested-scan checkpoint over the raw solver steps (endpoint-only, +explicit noise; verified against the library solve before timing). A chunked +win motivates a real checkpointing option in the package. + +Run on a GPU node (or CPU for a smoke test): + + python -m benchmarks.gpu_trajectories --batch 1000 10000 --n-steps 31 127 1024 + python -m benchmarks.gpu_trajectories --drifts ou mlp --modes grad --remat chunked + +Writes .json (raw records) and .md (summary table). +""" + +import argparse +import itertools +import json +import time +from pathlib import Path + +import numpy as np + +HIDDEN = 32 +SIGMA = 0.1 +THETA = 0.7 +TARGET_SECONDS_PER_REPEAT = 0.05 +MAX_ITERS = 1000 +POLICY_DIR = Path(__file__).parent / "policies" +GROWTH_DRIFTS = { + "neoclassical": "neoclassical_growth", + "stochastic-growth": "stochastic_growth", +} +GROWTH_LOG_Z_VARIANCE = 0.012564 +GROWTH_LOG_K_VARIANCE = 0.01 + + +def parse_args(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--solvers", nargs="+", default=["sra1", "em", "rk4"]) + parser.add_argument("--batch", nargs="+", type=int, default=[100, 1000, 10000]) + parser.add_argument("--dim", nargs="+", type=int, default=[1, 8]) + parser.add_argument( + "--n-steps", nargs="+", type=int, default=[31, 127, 1024], dest="n_steps" + ) + parser.add_argument("--dtypes", nargs="+", default=["float32"]) + parser.add_argument( + "--drifts", + nargs="+", + default=["ou"], + choices=["ou", "mlp", "neoclassical", "stochastic-growth"], + help="ou/mlp are synthetic; neoclassical (RK4 only) and " + "stochastic-growth (SDE solvers only) load trained kernels policies " + "from benchmarks/policies/ (see export_growth_policies.py)", + ) + parser.add_argument( + "--noise-modes", + nargs="+", + default=["explicit"], + choices=["explicit", "key"], + dest="noise_modes", + ) + parser.add_argument("--save", nargs="+", default=["t1"], choices=["t1", "steps"]) + parser.add_argument("--modes", nargs="+", default=["primal", "grad"]) + parser.add_argument( + "--remat", + default="none", + choices=["none", "solve", "chunked"], + help="reverse-mode memory policy: plain scan, jax.checkpoint around " + "the per-trajectory solve, or a sqrt(n)-chunked nested scan", + ) + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument( + "--matmul-precision", + default=None, + choices=["default", "tensorfloat32", "highest"], + dest="matmul_precision", + ) + parser.add_argument( + "--scan-unroll", + type=int, + default=1, + dest="scan_unroll", + help="forwarded to the solvers' unroll= argument", + ) + parser.add_argument( + "--output", + default="benchmarks/results/gpu_trajectories", + help="prefix for the .json and .md outputs", + ) + return parser.parse_args() + + +def unit_softplus(raw): + log_two = jnp.log(jnp.asarray(2.0, dtype=raw.dtype)) + return jax.nn.softplus(2.0 * log_two * raw) / log_two + + +def load_growth_model(drift_kind, dtype): + """Build (drift, diffusion, weights) from an exported kernels policy. + + The drift evaluates the trained investment-policy MLP at the (log z, + log k) state, exactly as the kernels package does; the weights pytree is + the differentiable `p` so grad mode measures the real training shape. + """ + name = GROWTH_DRIFTS[drift_kind] + metadata = json.loads((POLICY_DIR / f"{name}.json").read_text()) + raw = np.load(POLICY_DIR / f"{name}.npz") + weights = {key: jnp.asarray(raw[key], dtype) for key in raw.files} + structural = metadata["structural"] + num_layers = metadata["num_layers"] + homothetic = metadata["homothetic"] + log_z_scale = metadata["log_z_scale"] + log_k_scale = metadata["log_k_scale"] + k_scale = metadata["k_scale"] + approximate = metadata["gelu_approximate"] + delta = structural["delta"] + + def policy(log_state, w): + log_z, log_k = log_state + if homothetic: + h = jnp.atleast_1d(log_k - log_z) + anchor = delta * jnp.exp(log_k) + else: + h = jnp.stack([log_z - log_z_scale, log_k - log_k_scale]) + anchor = delta * k_scale + for index in range(num_layers): + h = jax.nn.gelu( + h @ w[f"layer{index}_kernel"] + w[f"layer{index}_bias"], + approximate=approximate, + ) + raw_head = h @ w["head_kernel"] + w["head_bias"] + return anchor * unit_softplus(raw_head)[0] + + if drift_kind == "neoclassical": + g = structural.get("g", 0.0) + + def drift(x, t, args, w): + log_k = x[1] + investment = policy(x, w) + return jnp.stack( + [ + jnp.asarray(g, x.dtype), + investment / jnp.exp(log_k) - delta, + ] + ) + + return drift, None, weights + + eta = structural["eta"] + log_z_bar = structural["log_z_bar"] + sigma_z = structural["sigma_z"] + + def drift(x, t, args, w): + log_z, log_k = x + investment = policy(x, w) + return jnp.stack( + [ + eta * (log_z_bar - log_z), + investment / jnp.exp(log_k) - delta, + ] + ) + + def diffusion(x, t, args, w): + return jnp.stack([jnp.asarray(sigma_z, x.dtype), jnp.zeros((), x.dtype)]) + + return drift, diffusion, weights + + +def growth_initial_states(batch, dtype): + key_z, key_k = jax.random.split(jax.random.key(7)) + log_z = GROWTH_LOG_Z_VARIANCE**0.5 * jax.random.normal(key_z, (batch,), dtype) + log_k = GROWTH_LOG_K_VARIANCE**0.5 * jax.random.normal(key_k, (batch,), dtype) + return jnp.stack([log_z, log_k], axis=1) + + +def ou_drift(x, t, args, p): + return -p * x + + +def mlp_drift(x, t, args, p): + return jnp.tanh(x @ p["w1"] + p["b1"]) @ p["w2"] + p["b2"] + + +def additive_diffusion(x, t, args, p): + return jnp.full_like(x, SIGMA) + + +def make_parameters(drift_kind, dim, dtype, key): + if drift_kind == "ou": + return jnp.asarray(THETA, dtype) + keys = jax.random.split(key, 4) + scale_1 = 1.0 / dim**0.5 + scale_2 = 1.0 / HIDDEN**0.5 + return { + "w1": scale_1 * jax.random.normal(keys[0], (dim, HIDDEN), dtype), + "b1": jnp.zeros((HIDDEN,), dtype), + "w2": scale_2 * jax.random.normal(keys[1], (HIDDEN, dim), dtype), + "b2": jnp.zeros((dim,), dtype), + } + + +def chunk_counts(n): + """Split n steps into (n_chunks, chunk) with n_chunks the largest divisor + of n at most sqrt(n); prime n degenerates to one whole-solve chunk.""" + best = 1 + divisor = 1 + while divisor * divisor <= n: + if n % divisor == 0: + best = divisor + divisor += 1 + return best, n // best + + +def chunked_sde_endpoint(solver, drift, diffusion, n, dt, dtype, unroll): + n_chunks, chunk = chunk_counts(n) + times = (jnp.arange(n, dtype=dtype) * dt).reshape(n_chunks, chunk) + + def one(x_0, p, noise): + def g_drift(x, t): + return drift(x, t, None, p) + + def g_diffusion(x, t): + return diffusion(x, t, None, p) + + def inner(x, inputs): + t, w = inputs + return solver.step(g_drift, g_diffusion, t, x, dt, w, identity), None + + def chunk_body(x, inputs): + x, _ = jax.lax.scan(inner, x, inputs, unroll=unroll) + return x, None + + noise_chunks = jax.tree.map( + lambda leaf: leaf.reshape((n_chunks, chunk) + leaf.shape[1:]), noise + ) + x, _ = jax.lax.scan(jax.checkpoint(chunk_body), x_0, (times, noise_chunks)) + return x + + return one + + +def chunked_rk4_endpoint(solver, drift, n, dt, dtype, unroll): + n_chunks, chunk = chunk_counts(n) + times = (jnp.arange(n, dtype=dtype) * dt).reshape(n_chunks, chunk) + + def one(x_0, p): + def g(x, t): + return drift(x, t, None, p) + + def inner(x, t): + x_1, _, _ = solver.step_fixed(g, t, x, dt, None, identity) + return x_1, None + + def chunk_body(x, ts): + x, _ = jax.lax.scan(inner, x, ts, unroll=unroll) + return x, None + + x, _ = jax.lax.scan(jax.checkpoint(chunk_body), x_0, times) + return x + + return one + + +def identity(x): + return x + + +def build_case(case, dtype, remat, unroll=1): + """Return (fn, inputs) where jit(fn)(*inputs) runs the whole ensemble.""" + save_at = SaveAt(steps=True) if case["save"] == "steps" else SaveAt(t_1=True) + n = case["n_steps"] + batch = case["batch"] + dim = case["dim"] + if case["drift"] in GROWTH_DRIFTS: + drift, diffusion, p = load_growth_model(case["drift"], dtype) + x_0s = growth_initial_states(batch, dtype) + else: + drift = ou_drift if case["drift"] == "ou" else mlp_drift + diffusion = additive_diffusion + p = make_parameters(case["drift"], dim, dtype, jax.random.key(0)) + x_0s = jnp.linspace(0.5, 2.0, batch * dim, dtype=dtype).reshape(batch, dim) + solver = {"sra1": SRA1(), "em": EulerMaruyama(), "rk4": RK4()}[case["solver"]] + dt = jnp.asarray(1.0 / n, dtype) + + if case["solver"] == "rk4": + if remat == "chunked": + one = chunked_rk4_endpoint(solver, drift, n, dt, dtype, unroll) + else: + + def one(x_0, p): + return solve_ode( + drift, + solver, + 0.0, + 1.0, + x_0, + p=p, + dt_0=1.0 / n, + max_steps=n, + save_at=save_at, + has_aux=False, + unroll=unroll, + ).xs + + if remat == "solve": + one = jax.checkpoint(one) + + if case["mode"] == "primal": + + def fn(x_0s, p): + return jax.vmap(lambda x_0: one(x_0, p))(x_0s) + + return fn, (x_0s, p) + + def loss(x_0s, p): + return jnp.sum(jax.vmap(lambda x_0: one(x_0, p))(x_0s) ** 2) + + return jax.grad(loss, argnums=(0, 1)), (x_0s, p) + + keys = jax.random.split(jax.random.key(1), batch) + if case["noise"] == "explicit": + noise = jax.vmap( + lambda k: solver.sample_noise(jnp.zeros((dim,), dtype), k, n, dt, dtype) + )(keys) + if remat == "chunked": + one = chunked_sde_endpoint(solver, drift, diffusion, n, dt, dtype, unroll) + else: + + def one(x_0, p, w): + return solve_sde( + drift, + diffusion, + solver, + 0.0, + 1.0, + x_0, + p=p, + noise=w, + n_steps=n, + save_at=save_at, + has_aux=False, + unroll=unroll, + ).xs + + if remat == "solve": + one = jax.checkpoint(one) + + if case["mode"] == "primal": + + def fn(x_0s, p, noise): + return jax.vmap(lambda x_0, w: one(x_0, p, w))(x_0s, noise) + + return fn, (x_0s, p, noise) + + def loss(x_0s, p, noise): + return jnp.sum(jax.vmap(lambda x_0, w: one(x_0, p, w))(x_0s, noise) ** 2) + + return jax.grad(loss, argnums=(0, 1, 2)), (x_0s, p, noise) + + def one(x_0, p, key): + return solve_sde( + drift, + diffusion, + solver, + 0.0, + 1.0, + x_0, + p=p, + key=key, + n_steps=n, + save_at=save_at, + has_aux=False, + unroll=unroll, + ).xs + + if remat == "solve": + one = jax.checkpoint(one) + + if case["mode"] == "primal": + + def fn(x_0s, p, keys): + return jax.vmap(lambda x_0, k: one(x_0, p, k))(x_0s, keys) + + return fn, (x_0s, p, keys) + + def loss(x_0s, p, keys): + return jnp.sum(jax.vmap(lambda x_0, k: one(x_0, p, k))(x_0s, keys) ** 2) + + return jax.grad(loss, argnums=(0, 1)), (x_0s, p, keys) + + +def verify_chunked(case, dtype): + """The chunked path bypasses the solve functions; require endpoint parity + before timing it so a silent mismatch cannot produce wrong conclusions.""" + primal = dict(case, mode="primal") + plain_fn, plain_inputs = build_case(primal, dtype, remat="none") + chunk_fn, chunk_inputs = build_case(primal, dtype, remat="chunked") + plain = jax.jit(plain_fn)(*plain_inputs) + chunked = jax.jit(chunk_fn)(*chunk_inputs) + tolerance = 200 * float(jnp.finfo(dtype).eps) + if not jnp.allclose(plain, chunked, atol=tolerance, rtol=tolerance): + gap = float(jnp.max(jnp.abs(plain - chunked))) + raise AssertionError(f"chunked endpoint mismatch for {case}: max |diff|={gap}") + + +def time_case(fn, inputs, repeats): + compiled = jax.jit(fn) + + def call(): + return jax.block_until_ready(compiled(*inputs)) + + start = time.perf_counter() + call() + compile_seconds = time.perf_counter() - start + single = time.perf_counter() + call() + single = time.perf_counter() - single + iters = max(1, min(MAX_ITERS, int(TARGET_SECONDS_PER_REPEAT / max(single, 1e-9)))) + best = float("inf") + for _ in range(repeats): + start = time.perf_counter() + for _ in range(iters): + call() + best = min(best, (time.perf_counter() - start) / iters) + return best, compile_seconds + + +def case_records(args, flush): + dtypes = {"float32": jnp.float32, "float64": jnp.float64} + records = [] + cases = [ + dict( + zip( + ("solver", "drift", "batch", "dim", "n_steps", "dtype", "save", "mode"), + values, + strict=True, + ) + ) + for values in itertools.product( + args.solvers, + args.drifts, + args.batch, + args.dim, + args.n_steps, + args.dtypes, + args.save, + args.modes, + ) + ] + for case in cases: + if case["drift"] in GROWTH_DRIFTS: + # Growth drifts bind to their model class: the neoclassical model + # is the deterministic ODE, stochastic growth is the SDE, and the + # state is fixed at (log z, log k). + if (case["drift"] == "neoclassical") != (case["solver"] == "rk4"): + continue + if case["dim"] != args.dim[0]: + continue + case = dict(case, dim=2) + noise_modes = ["none"] if case["solver"] == "rk4" else args.noise_modes + for noise_mode in noise_modes: + case = dict(case, noise=noise_mode) + if args.remat == "chunked" and ( + case["save"] != "t1" or noise_mode == "key" + ): + print(f"skip {case}: chunked remat is endpoint/explicit-noise only") + continue + dtype = dtypes[case["dtype"]] + try: + if args.remat == "chunked": + verify_chunked(case, dtype) + fn, inputs = build_case(case, dtype, args.remat, args.scan_unroll) + seconds, compile_seconds = time_case(fn, inputs, args.repeats) + except Exception as error: + # OOM is a measurement (the memory wall), not a crash; anything + # else still fails the run loudly. + if "RESOURCE_EXHAUSTED" not in str(error): + raise + record = dict(case, remat=args.remat, oom=True) + records.append(record) + flush(records) + print( + f"{case['solver']:>4} {case['drift']:>3} " + f"B={case['batch']:<6} d={case['dim']:<3} " + f"n={case['n_steps']:<5} {case['dtype']} " + f"noise={noise_mode:<8} save={case['save']:<5} " + f"{case['mode']:<6} OOM", + flush=True, + ) + continue + record = dict( + case, + remat=args.remat, + oom=False, + seconds_per_call=seconds, + trajectories_per_second=case["batch"] / seconds, + compile_seconds=compile_seconds, + ) + records.append(record) + flush(records) + print( + f"{case['solver']:>4} {case['drift']:>3} B={case['batch']:<6} " + f"d={case['dim']:<3} n={case['n_steps']:<5} {case['dtype']} " + f"noise={noise_mode:<8} save={case['save']:<5} " + f"{case['mode']:<6} {seconds * 1e3:9.3f} ms/call " + f"({record['trajectories_per_second']:.3g} traj/s)", + flush=True, + ) + return records + + +def write_outputs(records, args, device): + output = Path(args.output) + output.parent.mkdir(parents=True, exist_ok=True) + header = { + "device": device, + "backend": jax.default_backend(), + "jax_version": jax.__version__, + "matmul_precision": args.matmul_precision, + "scan_unroll": args.scan_unroll, + "remat": args.remat, + } + output.with_suffix(".json").write_text( + json.dumps({"config": header, "records": records}, indent=2) + "\n" + ) + lines = [ + f"# gpu_trajectories — {device} ({jax.default_backend()}, " + f"jax {jax.__version__})", + "", + f"matmul_precision={args.matmul_precision}, " + f"scan_unroll={args.scan_unroll}, remat={args.remat}", + "", + "| solver | drift | B | d | n_steps | dtype | noise | save | mode " + "| ms/call | traj/s | compile s |", + "|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:|", + ] + for r in records: + if r.get("oom"): + timing = "| OOM | — | — |" + else: + timing = ( + f"| {r['seconds_per_call'] * 1e3:.3f} " + f"| {r['trajectories_per_second']:.3g} " + f"| {r['compile_seconds']:.2f} |" + ) + lines.append( + f"| {r['solver']} | {r['drift']} | {r['batch']} | {r['dim']} " + f"| {r['n_steps']} | {r['dtype']} | {r['noise']} | {r['save']} " + f"| {r['mode']} {timing}" + ) + output.with_suffix(".md").write_text("\n".join(lines) + "\n") + + +def main(): + args = parse_args() + global jax, jnp, SRA1, EulerMaruyama, RK4, SaveAt, solve_ode, solve_sde + import jax + import jax.numpy as jnp + + if "float64" in args.dtypes: + jax.config.update("jax_enable_x64", True) + if args.matmul_precision is not None: + jax.config.update("jax_default_matmul_precision", args.matmul_precision) + from tinydiffeq import RK4, SRA1, EulerMaruyama, SaveAt, solve_ode, solve_sde + + device = jax.devices()[0].device_kind + print(f"device: {device} ({jax.default_backend()}), jax {jax.__version__}") + records = case_records( + args, flush=lambda records: write_outputs(records, args, device) + ) + write_outputs(records, args, device) + print(f"done: {Path(args.output).with_suffix('.json')}") + + +if __name__ == "__main__": + main() diff --git a/benchmarks/policies/neoclassical_growth.json b/benchmarks/policies/neoclassical_growth.json new file mode 100644 index 0000000..97c88c3 --- /dev/null +++ b/benchmarks/policies/neoclassical_growth.json @@ -0,0 +1,21 @@ +{ + "num_layers": 3, + "homothetic": true, + "log_z_scale": 0.0, + "log_k_scale": 0.0, + "k_scale": 1.0, + "structural": { + "a": 0.3333333333333333, + "delta": 0.1, + "rho_hat": 0.041817831615435494, + "sigma_crra": 2.0, + "g": 0.0, + "z_0": [ + 1.0 + ], + "k_0": [ + 1.0 + ] + }, + "gelu_approximate": true +} diff --git a/benchmarks/policies/neoclassical_growth.npz b/benchmarks/policies/neoclassical_growth.npz new file mode 100644 index 0000000..8bc0d9c Binary files /dev/null and b/benchmarks/policies/neoclassical_growth.npz differ diff --git a/benchmarks/policies/stochastic_growth.json b/benchmarks/policies/stochastic_growth.json new file mode 100644 index 0000000..6937465 --- /dev/null +++ b/benchmarks/policies/stochastic_growth.json @@ -0,0 +1,53 @@ +{ + "num_layers": 3, + "homothetic": false, + "log_z_scale": 0.0, + "log_k_scale": 0.0, + "k_scale": 1.0, + "structural": { + "a": 0.3333333333333333, + "delta": 0.1, + "rho_hat": 0.041817831615435494, + "sigma_crra": 2.0, + "eta": 0.2051731775502023, + "sigma_z": 0.07180274153740325, + "log_z_bar": 0.0, + "z_0": [ + 1.2296024610060845, + 1.0405861133091159, + 0.929477223262021, + 0.9996573515960677, + 0.8302528952965452, + 1.2036043826999496, + 1.0892246228947389, + 0.9565238483903568, + 1.1067454672391177, + 1.0770675044786562, + 0.8230255876255208, + 0.9266844557879376, + 0.9456675135848382, + 0.9276125419113377, + 1.056822169707839, + 0.9585549936963458 + ], + "k_0": [ + 0.9220961859877157, + 1.249409084056467, + 0.9329877067724575, + 0.9241332002131825, + 1.0010051935754307, + 1.077808522260105, + 0.8799155356142421, + 1.0492419562499204, + 0.9319484982916157, + 1.1321340981748897, + 0.9658447901861608, + 1.0576707064445756, + 1.123999157704123, + 0.8562897707666826, + 1.0412593731933835, + 0.9356658429406364 + ] + }, + "gelu_approximate": true +} diff --git a/benchmarks/policies/stochastic_growth.npz b/benchmarks/policies/stochastic_growth.npz new file mode 100644 index 0000000..4f26a5e Binary files /dev/null and b/benchmarks/policies/stochastic_growth.npz differ diff --git a/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md b/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md new file mode 100644 index 0000000..48d8623 --- /dev/null +++ b/benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md @@ -0,0 +1,93 @@ +# Vulcan L40S: fixed-step SDE/ODE trajectory ensembles (2026-08-04) + +Environment: NVIDIA L40S (48 GB), jax 0.9.1 + CUDA-12 plugin (Alliance +wheelhouse), Python 3.11, float32, TF32 matmul default. Workload: +`benchmarks/gpu_trajectories.py` with the **trained kernels investment +policies** (`benchmarks/policies/`, exported by `export_growth_policies.py`) — +the neoclassical growth ODE under RK4 and the stochastic growth SDE +(additive OU noise in log z) under SRA1/EulerMaruyama. Per-trajectory `x_0` +and explicit presampled `noise=`, vmapped and jitted; grad mode is +`jit(grad)` of a scalar loss with respect to `(x_0, policy weights, noise)`. +Raw records: `vulcan-l40s-*.json` in this directory. + +## Headline numbers (kernels scale: B ≤ 32, `SaveAt(steps=True)`) + +Milliseconds per whole-ensemble call, `unroll=1` → `unroll=8`: + +| model / solver | B | n_steps | primal | grad | +|---|---:|---:|---:|---:| +| stochastic growth / SRA1 | 32 | 63 | 1.35 → 1.09 | 8.03 → 3.43 | +| stochastic growth / SRA1 | 32 | 255 | 5.14 → 4.07 | 32.0 → 13.3 | +| stochastic growth / EM | 32 | 255 | 3.07 → 2.19 | 18.7 → 7.1 | +| neoclassical / RK4 | 32 | 63 | 2.11 → 1.62 | 14.3 → 6.0 | +| neoclassical / RK4 | 32 | 255 | 8.25 → 6.13 | 57.1 → 25.1 | + +Cost is set almost entirely by `n_steps` (sequential scan latency), not by +batch: times are flat in B from 1 to 32 — and in primal mode flat to +**B = 1024** (SRA1 n=255: 4.49 ms at B=16 vs 5.22 ms at B=1024). Forward +ensembles of hundreds of trajectories are effectively free; the grad-mode +per-LM-iteration cost at the kernels training shape (B=16, n=63, steps +saved, SRA1) is ~8 ms rolled and **~3.4 ms with `unroll=8`**. + +## Findings + +1. **`lax.scan` unroll is the one big lever: 2.0–3.1× faster grad, 1.2–1.7× + primal** at B ≤ 32 (per-iteration dispatch overhead dominates these small + kernels; the scan bodies are pure math — matmuls, elementwise, and their + transposes; no branches or bookkeeping to strip). Most of the gain is at + `unroll=4`; 8 adds a little more; compile time grows with the unroll + factor. Promoted to a real `unroll=` argument on `solve_sde` and + fixed-step `solve_ode`; `vulcan-l40s-g-unroll8-promoted` confirms the + argument reproduces the measured numbers. +2. **Explicit `noise=` is free**: identical timing to in-solve `key=` draws + in every configuration, so first-class noise costs nothing. +3. **`SaveAt(steps=True)` is free at this scale** (≤ 5% over endpoint-only), + so the collocation rollout shape carries no memory-traffic penalty. +4. **`jax_default_matmul_precision="highest"` costs 15–35%** on these + width-32 policy matmuls versus the TF32 default. Use it for + reproducibility when needed, knowingly. +5. **CUDA-graph command-buffer capture of the while loop + (`--xla_gpu_enable_command_buffer=...,WHILE`) had zero effect** — timings + identical to baseline. Unroll is the only dispatch-overhead remedy that + worked. +6. **The reverse-mode memory wall is real but far from kernels scale**: the + plain-scan grad tape OOMs the 48 GB card at B=10,000 × n=1024 (SRA1, + ~45 GiB requested; `vulcan-l40s-a-baseline` records the failure). At + B ≤ 32 the tape is tens of MB. The benchmark's `--remat chunked` mode + (√n nested-scan checkpoint, endpoint-parity-verified) is the demonstrated + remedy if ensembles ever grow to B·n ≳ 10⁷ in grad mode; a package + `checkpoint_every=` option was considered and deliberately not added. + +## Follow-up: vmapped adaptive solves (CPU measurement, same session) + +A batched-predicate `lax.cond` lowers under `vmap` to a both-branches +select, so vmapped **adaptive** solves used to execute every `max_steps` +attempt slot: a B=32 Tsit5 solve of the neoclassical policy model that +needs 8 attempts took 2.6 ms at `max_steps=64` but 47.5 ms at +`max_steps=1024` (CPU). Gating the skip conds on a scalar `unvmap_all` +predicate (batching rule reduces over the batch axis; the diffrax +`unvmap_any` trick) makes the vmapped primal budget-invariant: 0.5 ms at +every budget, a 95× win at `max_steps=1024`. Reverse mode improved 3.9× +(129 → 33 ms) but still scales with the budget through the scan's stacked +per-slot residuals — keep `max_steps` realistic when differentiating +adaptive solves. + +Validated on the L40S (same probe, pre-gate vs gated, ms per call): + +| B | max_steps | primal before → after | grad before → after | +|---:|---:|---|---| +| 32 | 64 | 2.05 → 0.87 (2.4×) | 9.0 → 5.7 (1.6×) | +| 32 | 256 | 7.82 → 1.16 (6.7×) | 36.2 → 15.4 (2.3×) | +| 32 | 1024 | 30.7 → 2.50 (12.3×) | 144.3 → 54.7 (2.6×) | +| 256 | 64 | 2.71 → 1.00 (2.7×) | — | +| 256 | 1024 | 40.9 → 2.59 (15.8×) | — | + +## Reproduce + +```bash +python -m benchmarks.gpu_trajectories \ + --drifts neoclassical stochastic-growth --solvers sra1 em rk4 \ + --batch 1 4 16 32 --dim 1 --n-steps 31 63 127 255 \ + --modes primal grad --save t1 steps --output +# knobs: --scan-unroll 8 | --matmul-precision highest | --noise-modes key +``` diff --git a/benchmarks/results/vulcan-l40s-a-baseline.json b/benchmarks/results/vulcan-l40s-a-baseline.json new file mode 100644 index 0000000..12381d9 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-a-baseline.json @@ -0,0 +1,1296 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 1, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006979090666087965, + "trajectories_per_second": 143285.14241247656, + "compile_seconds": 0.7613029390340671 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004477978091348301, + "trajectories_per_second": 22331.507202593395, + "compile_seconds": 2.764506147010252 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013014882979159419, + "trajectories_per_second": 76835.11266304033, + "compile_seconds": 0.44207365496549755 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008958171610720456, + "trajectories_per_second": 11162.992220457982, + "compile_seconds": 1.783474794938229 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002520166680608925, + "trajectories_per_second": 39679.915129993664, + "compile_seconds": 0.471507192007266 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01790090400027111, + "trajectories_per_second": 5586.310054424373, + "compile_seconds": 1.7965833969647065 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.019478568516205996, + "trajectories_per_second": 5133.847485599411, + "compile_seconds": 0.45416608604136854 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.15820773702580482, + "trajectories_per_second": 632.0803386732551, + "compile_seconds": 1.9938357799546793 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007358905841381504, + "trajectories_per_second": 1358897.6697822073, + "compile_seconds": 1.2876181559404358 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005326191778294742, + "trajectories_per_second": 187751.40693866726, + "compile_seconds": 4.128417236963287 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001382418572237449, + "trajectories_per_second": 723369.9113152804, + "compile_seconds": 0.4274732309859246 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01103723022970371, + "trajectories_per_second": 90602.44093747102, + "compile_seconds": 1.6469531509792432 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00268453383533698, + "trajectories_per_second": 372504.1520568034, + "compile_seconds": 0.45855261699762195 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.02228194155031815, + "trajectories_per_second": 44879.39247761475, + "compile_seconds": 1.7092685449169949 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.02078921248903498, + "trajectories_per_second": 48101.87016595448, + "compile_seconds": 0.4331063599092886 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.18087553698569536, + "trajectories_per_second": 5528.66361402474, + "compile_seconds": 1.86476408992894 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009766825313242723, + "trajectories_per_second": 10238741.535021741, + "compile_seconds": 0.5928605200024322 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.011672457243548706, + "trajectories_per_second": 856717.638055769, + "compile_seconds": 2.0495360889472067 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0018726597683360944, + "trajectories_per_second": 5339998.31100406, + "compile_seconds": 0.44451983796898276 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.023683579463977367, + "trajectories_per_second": 422233.4725715748, + "compile_seconds": 1.7391729099908844 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003682566538023261, + "trajectories_per_second": 2715497.438199129, + "compile_seconds": 0.4114024490118027 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.04794425202999264, + "trajectories_per_second": 208575.5763536423, + "compile_seconds": 1.8183223910164088 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.029412479023449123, + "trajectories_per_second": 339991.74268947175, + "compile_seconds": 0.4274300360120833 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": true + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002987449995998759, + "trajectories_per_second": 33473363.615770977, + "compile_seconds": 0.6725022579776123 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.09317554696463048, + "trajectories_per_second": 1073242.961889562, + "compile_seconds": 4.375617903075181 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00600924136233516, + "trajectories_per_second": 16641035.693254383, + "compile_seconds": 0.4181409140583128 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.19127922097686678, + "trajectories_per_second": 522795.9393043218, + "compile_seconds": 2.0444391909986734 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.012165952502982691, + "trajectories_per_second": 8219660.562991948, + "compile_seconds": 0.49205413402523845 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": true + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.09763287496753037, + "trajectories_per_second": 1024245.1636629247, + "compile_seconds": 0.571541424957104 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": true + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0004140019734305414, + "trajectories_per_second": 241544.74233871582, + "compile_seconds": 0.26524597394745797 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0026054690550903186, + "trajectories_per_second": 38380.80510095849, + "compile_seconds": 1.145113329985179 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007423295588153651, + "trajectories_per_second": 134711.0576595972, + "compile_seconds": 0.24592959508299828 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0051357642255930435, + "trajectories_per_second": 19471.298838383234, + "compile_seconds": 1.1722897390136495 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0014299512635368634, + "trajectories_per_second": 69932.45332897463, + "compile_seconds": 0.2810975379543379 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01015592849580571, + "trajectories_per_second": 9846.465543874096, + "compile_seconds": 1.1618733649374917 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01042707500164397, + "trajectories_per_second": 9590.417253566666, + "compile_seconds": 0.2653352909255773 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.08733745594508946, + "trajectories_per_second": 1144.9841184161776, + "compile_seconds": 1.2963226590072736 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00045979170171522545, + "trajectories_per_second": 2174897.885867796, + "compile_seconds": 0.27072585409041494 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0028724446413381133, + "trajectories_per_second": 348135.5169073529, + "compile_seconds": 1.2190662970533594 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008201496548612009, + "trajectories_per_second": 1219289.6675293196, + "compile_seconds": 0.2859465579967946 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00596189200587105, + "trajectories_per_second": 167731.98827070955, + "compile_seconds": 1.144728722050786 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0015455124666914344, + "trajectories_per_second": 647034.573678177, + "compile_seconds": 0.28703883197158575 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.012067058734828606, + "trajectories_per_second": 82870.23557064036, + "compile_seconds": 1.1576002539368346 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.011439821013482288, + "trajectories_per_second": 87413.95506288603, + "compile_seconds": 0.2860855530016124 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.09809419699013233, + "trajectories_per_second": 10194.282951320696, + "compile_seconds": 1.2522407149663195 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0005777520452474439, + "trajectories_per_second": 17308463.175958343, + "compile_seconds": 0.24972683808300644 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006051044125342742, + "trajectories_per_second": 1652607.35054276, + "compile_seconds": 1.1796769490465522 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001048030680976808, + "trajectories_per_second": 9541705.392326478, + "compile_seconds": 0.2818208369426429 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.012286034994758666, + "trajectories_per_second": 813932.2413021036, + "compile_seconds": 1.1883608979405835 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0020266218293134284, + "trajectories_per_second": 4934319.691694905, + "compile_seconds": 0.29116004006937146 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.024754316080361605, + "trajectories_per_second": 403969.956897065, + "compile_seconds": 1.2487655740696937 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.015777506011848647, + "trajectories_per_second": 633813.7340901765, + "compile_seconds": 0.2851269660750404 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 10000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.19926861603744328, + "trajectories_per_second": 50183.51709795066, + "compile_seconds": 1.4213094549486414 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0015345037681981921, + "trajectories_per_second": 65167647.07421969, + "compile_seconds": 0.2797583380015567 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.04465881793294102, + "trajectories_per_second": 2239199.4376151743, + "compile_seconds": 1.4061723200138658 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00315982960164547, + "trajectories_per_second": 31647276.15309552, + "compile_seconds": 0.3583576310193166 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.09023767802864313, + "trajectories_per_second": 1108184.5431379355, + "compile_seconds": 1.484362359973602 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006186905716146741, + "trajectories_per_second": 16163168.56728841, + "compile_seconds": 0.3430421269731596 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.18328238592948765, + "trajectories_per_second": 545606.1666420688, + "compile_seconds": 1.5297465489711612 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.049338942975737154, + "trajectories_per_second": 2026796.5620823263, + "compile_seconds": 0.3346744281006977 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 100000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": true + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001035444549329224, + "trajectories_per_second": 96576.87614926502, + "compile_seconds": 0.5566155769629404 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007564925181213766, + "trajectories_per_second": 13218.901390899855, + "compile_seconds": 2.647828501998447 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001985040695770927, + "trajectories_per_second": 50376.800945717216, + "compile_seconds": 0.4893043019110337 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01516536200263848, + "trajectories_per_second": 6593.973818930398, + "compile_seconds": 2.6690926519222558 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003880432087250731, + "trajectories_per_second": 25770.32602337065, + "compile_seconds": 0.47296555899083614 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.031245502992533147, + "trajectories_per_second": 3200.460559841119, + "compile_seconds": 2.734248523018323 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.030213839025236666, + "trajectories_per_second": 3309.741602729569, + "compile_seconds": 0.4995520890224725 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 100, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.26694091293029487, + "trajectories_per_second": 374.614737404875, + "compile_seconds": 3.0433065260294825 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011527200516026753, + "trajectories_per_second": 867513.3208706293, + "compile_seconds": 0.49437285808380693 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.010639087238814682, + "trajectories_per_second": 93993.02567533149, + "compile_seconds": 2.76271853398066 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021989913181063125, + "trajectories_per_second": 454753.95549135766, + "compile_seconds": 0.5211080589797348 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.021824156516231596, + "trajectories_per_second": 45820.785754366065, + "compile_seconds": 2.7108210891019553 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004315479097633876, + "trajectories_per_second": 231723.9818281793, + "compile_seconds": 0.49992516089696437 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.04420895501971245, + "trajectories_per_second": 22619.85155618601, + "compile_seconds": 2.7040379130048677 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03372074698563665, + "trajectories_per_second": 29655.33356737174, + "compile_seconds": 0.5715859179617837 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1000, + "dim": 2, + "n_steps": 1024, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.35841227206401527, + "trajectories_per_second": 2790.0830354977134, + "compile_seconds": 3.108547538984567 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 10000, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001657568448191059, + "trajectories_per_second": 6032933.367495756, + "compile_seconds": 0.48741554596927017 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-a-baseline.md b/benchmarks/results/vulcan-l40s-a-baseline.md new file mode 100644 index 0000000..c827ecf --- /dev/null +++ b/benchmarks/results/vulcan-l40s-a-baseline.md @@ -0,0 +1,87 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=1, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | primal | 0.698 | 1.43e+05 | 0.76 | +| sra1 | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | grad | 4.478 | 2.23e+04 | 2.76 | +| sra1 | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | primal | 1.301 | 7.68e+04 | 0.44 | +| sra1 | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | grad | 8.958 | 1.12e+04 | 1.78 | +| sra1 | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | primal | 2.520 | 3.97e+04 | 0.47 | +| sra1 | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | grad | 17.901 | 5.59e+03 | 1.80 | +| sra1 | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | primal | 19.479 | 5.13e+03 | 0.45 | +| sra1 | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | grad | 158.208 | 632 | 1.99 | +| sra1 | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | primal | 0.736 | 1.36e+06 | 1.29 | +| sra1 | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | grad | 5.326 | 1.88e+05 | 4.13 | +| sra1 | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | primal | 1.382 | 7.23e+05 | 0.43 | +| sra1 | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | grad | 11.037 | 9.06e+04 | 1.65 | +| sra1 | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | primal | 2.685 | 3.73e+05 | 0.46 | +| sra1 | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | grad | 22.282 | 4.49e+04 | 1.71 | +| sra1 | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | primal | 20.789 | 4.81e+04 | 0.43 | +| sra1 | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | grad | 180.876 | 5.53e+03 | 1.86 | +| sra1 | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | primal | 0.977 | 1.02e+07 | 0.59 | +| sra1 | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | grad | 11.672 | 8.57e+05 | 2.05 | +| sra1 | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | primal | 1.873 | 5.34e+06 | 0.44 | +| sra1 | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | grad | 23.684 | 4.22e+05 | 1.74 | +| sra1 | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | primal | 3.683 | 2.72e+06 | 0.41 | +| sra1 | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | grad | 47.944 | 2.09e+05 | 1.82 | +| sra1 | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | primal | 29.412 | 3.4e+05 | 0.43 | +| sra1 | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | grad | OOM | — | — | +| sra1 | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | primal | 2.987 | 3.35e+07 | 0.67 | +| sra1 | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | grad | 93.176 | 1.07e+06 | 4.38 | +| sra1 | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | primal | 6.009 | 1.66e+07 | 0.42 | +| sra1 | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | grad | 191.279 | 5.23e+05 | 2.04 | +| sra1 | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | primal | 12.166 | 8.22e+06 | 0.49 | +| sra1 | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | grad | OOM | — | — | +| sra1 | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | primal | 97.633 | 1.02e+06 | 0.57 | +| sra1 | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | grad | OOM | — | — | +| em | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | primal | 0.414 | 2.42e+05 | 0.27 | +| em | stochastic-growth | 100 | 2 | 31 | float32 | explicit | t1 | grad | 2.605 | 3.84e+04 | 1.15 | +| em | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | primal | 0.742 | 1.35e+05 | 0.25 | +| em | stochastic-growth | 100 | 2 | 63 | float32 | explicit | t1 | grad | 5.136 | 1.95e+04 | 1.17 | +| em | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | primal | 1.430 | 6.99e+04 | 0.28 | +| em | stochastic-growth | 100 | 2 | 127 | float32 | explicit | t1 | grad | 10.156 | 9.85e+03 | 1.16 | +| em | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | primal | 10.427 | 9.59e+03 | 0.27 | +| em | stochastic-growth | 100 | 2 | 1024 | float32 | explicit | t1 | grad | 87.337 | 1.14e+03 | 1.30 | +| em | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | primal | 0.460 | 2.17e+06 | 0.27 | +| em | stochastic-growth | 1000 | 2 | 31 | float32 | explicit | t1 | grad | 2.872 | 3.48e+05 | 1.22 | +| em | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | primal | 0.820 | 1.22e+06 | 0.29 | +| em | stochastic-growth | 1000 | 2 | 63 | float32 | explicit | t1 | grad | 5.962 | 1.68e+05 | 1.14 | +| em | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | primal | 1.546 | 6.47e+05 | 0.29 | +| em | stochastic-growth | 1000 | 2 | 127 | float32 | explicit | t1 | grad | 12.067 | 8.29e+04 | 1.16 | +| em | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | primal | 11.440 | 8.74e+04 | 0.29 | +| em | stochastic-growth | 1000 | 2 | 1024 | float32 | explicit | t1 | grad | 98.094 | 1.02e+04 | 1.25 | +| em | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | primal | 0.578 | 1.73e+07 | 0.25 | +| em | stochastic-growth | 10000 | 2 | 31 | float32 | explicit | t1 | grad | 6.051 | 1.65e+06 | 1.18 | +| em | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | primal | 1.048 | 9.54e+06 | 0.28 | +| em | stochastic-growth | 10000 | 2 | 63 | float32 | explicit | t1 | grad | 12.286 | 8.14e+05 | 1.19 | +| em | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | primal | 2.027 | 4.93e+06 | 0.29 | +| em | stochastic-growth | 10000 | 2 | 127 | float32 | explicit | t1 | grad | 24.754 | 4.04e+05 | 1.25 | +| em | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | primal | 15.778 | 6.34e+05 | 0.29 | +| em | stochastic-growth | 10000 | 2 | 1024 | float32 | explicit | t1 | grad | 199.269 | 5.02e+04 | 1.42 | +| em | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | primal | 1.535 | 6.52e+07 | 0.28 | +| em | stochastic-growth | 100000 | 2 | 31 | float32 | explicit | t1 | grad | 44.659 | 2.24e+06 | 1.41 | +| em | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | primal | 3.160 | 3.16e+07 | 0.36 | +| em | stochastic-growth | 100000 | 2 | 63 | float32 | explicit | t1 | grad | 90.238 | 1.11e+06 | 1.48 | +| em | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | primal | 6.187 | 1.62e+07 | 0.34 | +| em | stochastic-growth | 100000 | 2 | 127 | float32 | explicit | t1 | grad | 183.282 | 5.46e+05 | 1.53 | +| em | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | primal | 49.339 | 2.03e+06 | 0.33 | +| em | stochastic-growth | 100000 | 2 | 1024 | float32 | explicit | t1 | grad | OOM | — | — | +| rk4 | neoclassical | 100 | 2 | 31 | float32 | none | t1 | primal | 1.035 | 9.66e+04 | 0.56 | +| rk4 | neoclassical | 100 | 2 | 31 | float32 | none | t1 | grad | 7.565 | 1.32e+04 | 2.65 | +| rk4 | neoclassical | 100 | 2 | 63 | float32 | none | t1 | primal | 1.985 | 5.04e+04 | 0.49 | +| rk4 | neoclassical | 100 | 2 | 63 | float32 | none | t1 | grad | 15.165 | 6.59e+03 | 2.67 | +| rk4 | neoclassical | 100 | 2 | 127 | float32 | none | t1 | primal | 3.880 | 2.58e+04 | 0.47 | +| rk4 | neoclassical | 100 | 2 | 127 | float32 | none | t1 | grad | 31.246 | 3.2e+03 | 2.73 | +| rk4 | neoclassical | 100 | 2 | 1024 | float32 | none | t1 | primal | 30.214 | 3.31e+03 | 0.50 | +| rk4 | neoclassical | 100 | 2 | 1024 | float32 | none | t1 | grad | 266.941 | 375 | 3.04 | +| rk4 | neoclassical | 1000 | 2 | 31 | float32 | none | t1 | primal | 1.153 | 8.68e+05 | 0.49 | +| rk4 | neoclassical | 1000 | 2 | 31 | float32 | none | t1 | grad | 10.639 | 9.4e+04 | 2.76 | +| rk4 | neoclassical | 1000 | 2 | 63 | float32 | none | t1 | primal | 2.199 | 4.55e+05 | 0.52 | +| rk4 | neoclassical | 1000 | 2 | 63 | float32 | none | t1 | grad | 21.824 | 4.58e+04 | 2.71 | +| rk4 | neoclassical | 1000 | 2 | 127 | float32 | none | t1 | primal | 4.315 | 2.32e+05 | 0.50 | +| rk4 | neoclassical | 1000 | 2 | 127 | float32 | none | t1 | grad | 44.209 | 2.26e+04 | 2.70 | +| rk4 | neoclassical | 1000 | 2 | 1024 | float32 | none | t1 | primal | 33.721 | 2.97e+04 | 0.57 | +| rk4 | neoclassical | 1000 | 2 | 1024 | float32 | none | t1 | grad | 358.412 | 2.79e+03 | 3.11 | +| rk4 | neoclassical | 10000 | 2 | 31 | float32 | none | t1 | primal | 1.658 | 6.03e+06 | 0.49 | diff --git a/benchmarks/results/vulcan-l40s-a-kernels-scale.json b/benchmarks/results/vulcan-l40s-a-kernels-scale.json new file mode 100644 index 0000000..0c17191 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-a-kernels-scale.json @@ -0,0 +1,3084 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 1, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00044518015159275196, + "trajectories_per_second": 2246.28163771954, + "compile_seconds": 0.46882447390817106 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002341831762654086, + "trajectories_per_second": 427.01615715838716, + "compile_seconds": 1.4274521339684725 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0005290067263508904, + "trajectories_per_second": 1890.3351322166357, + "compile_seconds": 0.33365265605971217 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0024705242766584787, + "trajectories_per_second": 404.7723835171357, + "compile_seconds": 1.284887518035248 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000821885879649299, + "trajectories_per_second": 1216.713931655211, + "compile_seconds": 0.3873317010002211 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004553293704520911, + "trajectories_per_second": 219.62123791995054, + "compile_seconds": 1.376192542957142 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009717659146796601, + "trajectories_per_second": 1029.0544100115378, + "compile_seconds": 0.34592496894765645 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004817036399617791, + "trajectories_per_second": 207.59652139629776, + "compile_seconds": 1.2965855040820315 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0015334397734653566, + "trajectories_per_second": 652.128643918073, + "compile_seconds": 0.3606032660463825 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00902233279775828, + "trajectories_per_second": 110.83608002671586, + "compile_seconds": 1.3879702399717644 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0018500350834801792, + "trajectories_per_second": 540.5302899006962, + "compile_seconds": 0.34839765599463135 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009488765196874738, + "trajectories_per_second": 105.38779063996276, + "compile_seconds": 1.3026674310676754 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003001927128934767, + "trajectories_per_second": 333.119345356944, + "compile_seconds": 0.3560068689985201 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01782353053567931, + "trajectories_per_second": 56.1056070231535, + "compile_seconds": 1.3728208880638704 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0035638541538411607, + "trajectories_per_second": 280.5950964413595, + "compile_seconds": 0.3729558620834723 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01892375253373757, + "trajectories_per_second": 52.84364177861573, + "compile_seconds": 1.2915878460044041 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006183858495205641, + "trajectories_per_second": 6468.453317780814, + "compile_seconds": 0.5545572639675811 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004215102452276783, + "trajectories_per_second": 948.9686301312568, + "compile_seconds": 1.6543556730030105 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006971375533699323, + "trajectories_per_second": 5737.7485700837315, + "compile_seconds": 0.37709420709870756 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004352877636186101, + "trajectories_per_second": 918.9323326590718, + "compile_seconds": 1.3398378629935905 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011719442831928887, + "trajectories_per_second": 3413.131543337752, + "compile_seconds": 0.3531762290513143 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008354840008541942, + "trajectories_per_second": 478.76440433454405, + "compile_seconds": 1.3873992669396102 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013069192043923099, + "trajectories_per_second": 3060.632965340743, + "compile_seconds": 0.3521754309767857 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008634187793359161, + "trajectories_per_second": 463.2746120111647, + "compile_seconds": 1.3849235730012879 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0022584000485949217, + "trajectories_per_second": 1771.1653887399737, + "compile_seconds": 0.3622201889520511 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.016870173043571413, + "trajectories_per_second": 237.10485895248414, + "compile_seconds": 1.391697199898772 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002554644420007734, + "trajectories_per_second": 1565.7756393306158, + "compile_seconds": 0.3469811760587618 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.017001587024424225, + "trajectories_per_second": 235.27215396148958, + "compile_seconds": 1.3942796170013025 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004364936189218001, + "trajectories_per_second": 916.3936943409518, + "compile_seconds": 0.34337302402127534 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03425101703032851, + "trajectories_per_second": 116.78485332152587, + "compile_seconds": 1.4128429669653997 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0050102312298905514, + "trajectories_per_second": 798.3663460752848, + "compile_seconds": 0.3538889540359378 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03452990192454308, + "trajectories_per_second": 115.84162644716027, + "compile_seconds": 1.4081455019768327 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006446458275119463, + "trajectories_per_second": 24819.83023415054, + "compile_seconds": 0.491259106900543 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003999261545356025, + "trajectories_per_second": 4000.7385909979635, + "compile_seconds": 1.5653145679971203 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007215623944913121, + "trajectories_per_second": 22174.104584926015, + "compile_seconds": 0.35081532900221646 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004105277084322138, + "trajectories_per_second": 3897.422676073012, + "compile_seconds": 1.3986358649563044 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001194223682055386, + "trajectories_per_second": 13397.825081196093, + "compile_seconds": 0.3312612830195576 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007826191159741333, + "trajectories_per_second": 2044.4172233238453, + "compile_seconds": 1.4205455440096557 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013475033638039322, + "trajectories_per_second": 11873.810804325438, + "compile_seconds": 0.35316630301531404 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008141389000229537, + "trajectories_per_second": 1965.2666147691625, + "compile_seconds": 1.3912285269470885 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002282388186791823, + "trajectories_per_second": 7010.201022153889, + "compile_seconds": 0.36521074501797557 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.015496470034122467, + "trajectories_per_second": 1032.4932042438559, + "compile_seconds": 1.3830204929690808 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002586590333117379, + "trajectories_per_second": 6185.749554208174, + "compile_seconds": 0.35045432997867465 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.016209229012019932, + "trajectories_per_second": 987.0919824832644, + "compile_seconds": 1.427961067063734 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004546631302218884, + "trajectories_per_second": 3519.0889554188284, + "compile_seconds": 0.343484070035629 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03083732898812741, + "trajectories_per_second": 518.8516815499848, + "compile_seconds": 1.4297372420551255 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00505759821842528, + "trajectories_per_second": 3163.5569511454223, + "compile_seconds": 0.3520263599930331 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03187646996229887, + "trajectories_per_second": 501.93763672463155, + "compile_seconds": 1.4496542459819466 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006498913063357274, + "trajectories_per_second": 49239.00302717562, + "compile_seconds": 0.6055959139484912 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003941296415481095, + "trajectories_per_second": 8119.155888480393, + "compile_seconds": 1.6412975949933752 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007238080008671834, + "trajectories_per_second": 44210.61933780959, + "compile_seconds": 0.407810817938298 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004049962997669354, + "trajectories_per_second": 7901.306757225967, + "compile_seconds": 1.5103212649701163 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0012194689998805926, + "trajectories_per_second": 26240.929456290698, + "compile_seconds": 0.43304388399701566 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007830712827853858, + "trajectories_per_second": 4086.47344162283, + "compile_seconds": 1.4790805460652336 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013507371259038337, + "trajectories_per_second": 23690.768089747653, + "compile_seconds": 0.4133943960769102 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008031683488904187, + "trajectories_per_second": 3984.220748266309, + "compile_seconds": 1.5542903270106763 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00238833566462355, + "trajectories_per_second": 13398.451680804192, + "compile_seconds": 0.4024871389847249 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.015544181650814911, + "trajectories_per_second": 2058.648098616525, + "compile_seconds": 1.527874932042323 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002630457443754292, + "trajectories_per_second": 12165.184453365779, + "compile_seconds": 0.4074179329909384 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0159813833112518, + "trajectories_per_second": 2002.3297969125229, + "compile_seconds": 1.5075764500070363 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00456721659284085, + "trajectories_per_second": 7006.455540155522, + "compile_seconds": 0.4330545160919428 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.031249569030478597, + "trajectories_per_second": 1024.014122204037, + "compile_seconds": 1.547785917064175 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00513628844378723, + "trajectories_per_second": 6230.1797008123, + "compile_seconds": 0.4154914850369096 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03203625499736518, + "trajectories_per_second": 998.8683134976869, + "compile_seconds": 1.5479462939547375 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00030696436657871997, + "trajectories_per_second": 3257.7071115632352, + "compile_seconds": 0.18483830604236573 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013746748898281818, + "trajectories_per_second": 727.4447270401428, + "compile_seconds": 0.7391808689571917 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0003924223752359727, + "trajectories_per_second": 2548.274673172438, + "compile_seconds": 0.22195479401852936 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0014773202728423657, + "trajectories_per_second": 676.901291062634, + "compile_seconds": 0.7877419149735942 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0005372288382008068, + "trajectories_per_second": 1861.4041706119608, + "compile_seconds": 0.19184806395787746 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0026184850010193056, + "trajectories_per_second": 381.90022078061435, + "compile_seconds": 0.7403571859467775 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006974619644661916, + "trajectories_per_second": 1433.7699415126651, + "compile_seconds": 0.22224036301486194 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0028671830043415815, + "trajectories_per_second": 348.7743888289543, + "compile_seconds": 0.7475595809519291 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00101398247749206, + "trajectories_per_second": 986.2103361720375, + "compile_seconds": 0.19123999297153205 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005112051782715652, + "trajectories_per_second": 195.6161718433874, + "compile_seconds": 0.74928293004632 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013029114411705557, + "trajectories_per_second": 767.5118725656324, + "compile_seconds": 0.2219563260441646 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005600566626526415, + "trajectories_per_second": 178.55336195156028, + "compile_seconds": 0.7629783579614013 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0019393865416835372, + "trajectories_per_second": 515.6269668304096, + "compile_seconds": 0.19201537396293133 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.010114373755641282, + "trajectories_per_second": 98.86919587505366, + "compile_seconds": 0.7507086748955771 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025219939480949608, + "trajectories_per_second": 396.511657276327, + "compile_seconds": 0.2333108299644664 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.010971961746690795, + "trajectories_per_second": 91.1414041615307, + "compile_seconds": 0.7984069479862228 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0003989394178491209, + "trajectories_per_second": 10026.585042826735, + "compile_seconds": 0.2265496940817684 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002401275502052158, + "trajectories_per_second": 1665.7813718507323, + "compile_seconds": 0.9170253330375999 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00047566794711066055, + "trajectories_per_second": 8409.227538448014, + "compile_seconds": 0.24313144909683615 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025512944428353673, + "trajectories_per_second": 1567.8315810364177, + "compile_seconds": 0.9244248060276732 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007076707728573996, + "trajectories_per_second": 5652.345911996605, + "compile_seconds": 0.21988887200132012 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0046742793987505134, + "trajectories_per_second": 855.7468774907303, + "compile_seconds": 0.9217888719867915 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000854283126748421, + "trajectories_per_second": 4682.288429627342, + "compile_seconds": 0.24327497300691903 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005006728784388138, + "trajectories_per_second": 798.9248413999784, + "compile_seconds": 0.9286532551050186 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013328823331517704, + "trajectories_per_second": 3001.0150937641206, + "compile_seconds": 0.22099014394916594 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009231478604488075, + "trajectories_per_second": 433.30003473715647, + "compile_seconds": 0.9346064699348062 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0016363451431971043, + "trajectories_per_second": 2444.4720703511043, + "compile_seconds": 0.2426328459987417 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00998435850488022, + "trajectories_per_second": 400.62663996338415, + "compile_seconds": 0.9612002939684317 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025957618398885977, + "trajectories_per_second": 1540.9734200313492, + "compile_seconds": 0.22377602802589536 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01833497901679948, + "trajectories_per_second": 218.16223494638243, + "compile_seconds": 0.9413322400068864 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0031615220631162327, + "trajectories_per_second": 1265.213375122646, + "compile_seconds": 0.24617118400055915 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01979129499522969, + "trajectories_per_second": 202.10905860198247, + "compile_seconds": 0.9616675300057977 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0003946929003347407, + "trajectories_per_second": 40537.84597197044, + "compile_seconds": 0.21592206100467592 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002193548297509551, + "trajectories_per_second": 7294.117944959603, + "compile_seconds": 0.9024808730464429 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00048250114121354924, + "trajectories_per_second": 33160.5433300283, + "compile_seconds": 0.23947630892507732 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002308497050156196, + "trajectories_per_second": 6930.916372155389, + "compile_seconds": 0.9371393930632621 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007099968506329096, + "trajectories_per_second": 22535.311228123315, + "compile_seconds": 0.21416236797813326 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004313542645170607, + "trajectories_per_second": 3709.2481322546832, + "compile_seconds": 0.9356112200766802 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008749259607296209, + "trajectories_per_second": 18287.261686299982, + "compile_seconds": 0.2405970450490713 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004538466199301183, + "trajectories_per_second": 3525.4201083316702, + "compile_seconds": 0.9290328299393877 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013390144723577097, + "trajectories_per_second": 11949.086683004645, + "compile_seconds": 0.2172669160645455 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00847135940566659, + "trajectories_per_second": 1888.7169383106818, + "compile_seconds": 0.9089790439466015 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001627106498926878, + "trajectories_per_second": 9833.406731859559, + "compile_seconds": 0.2721616820199415 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008918875199742615, + "trajectories_per_second": 1793.9481875990075, + "compile_seconds": 0.920522270957008 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025571096145237484, + "trajectories_per_second": 6257.064581480578, + "compile_seconds": 0.21738911396823823 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.016777171520516276, + "trajectories_per_second": 953.6768447788772, + "compile_seconds": 0.929095889092423 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003161054065761467, + "trajectories_per_second": 5061.602765135166, + "compile_seconds": 0.2428266960196197 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.017750939005054533, + "trajectories_per_second": 901.3607671934449, + "compile_seconds": 0.9496175439562649 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00038076739785877073, + "trajectories_per_second": 84040.80858800055, + "compile_seconds": 0.2364281170303002 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002288160718134826, + "trajectories_per_second": 13985.031622291164, + "compile_seconds": 0.9443110490683466 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00046480763727102916, + "trajectories_per_second": 68845.68461025698, + "compile_seconds": 0.2578090919414535 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0024021743171799338, + "trajectories_per_second": 13321.264727185515, + "compile_seconds": 0.9402767500141636 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006878971490853776, + "trajectories_per_second": 46518.582091155535, + "compile_seconds": 0.23492574703413993 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004524712641300125, + "trajectories_per_second": 7072.272326846631, + "compile_seconds": 0.924483462003991 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000842625876040464, + "trajectories_per_second": 37976.52185851378, + "compile_seconds": 0.25491039804182947 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004737628798466176, + "trajectories_per_second": 6754.4337813802795, + "compile_seconds": 0.970670091919601 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0012932961076699398, + "trajectories_per_second": 24742.98021174179, + "compile_seconds": 0.23251260293181986 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00888781079556793, + "trajectories_per_second": 3600.436680758032, + "compile_seconds": 0.9200960400048643 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001609387603821233, + "trajectories_per_second": 19883.339429246957, + "compile_seconds": 0.25817957904655486 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009383920207619667, + "trajectories_per_second": 3410.0886721112847, + "compile_seconds": 0.9416598289972171 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002494488738624281, + "trajectories_per_second": 12828.280001635969, + "compile_seconds": 0.2375009689712897 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.017782647511921823, + "trajectories_per_second": 1799.5070744413392, + "compile_seconds": 0.9460055059753358 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0030722367446287535, + "trajectories_per_second": 10415.863964893388, + "compile_seconds": 0.2579932389780879 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01872129802359268, + "trajectories_per_second": 1709.2831896417347, + "compile_seconds": 0.9482306479476392 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007537903852368656, + "trajectories_per_second": 1326.6287546049919, + "compile_seconds": 0.5394193780375645 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00450661373113028, + "trajectories_per_second": 221.8960975271327, + "compile_seconds": 2.2642879619961604 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009075410945235559, + "trajectories_per_second": 1101.8784780484057, + "compile_seconds": 0.7249034650158137 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004634865303523838, + "trajectories_per_second": 215.75600033936493, + "compile_seconds": 2.511369476094842 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0014247972079936195, + "trajectories_per_second": 701.85426697192, + "compile_seconds": 0.4596191640011966 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00899195959791541, + "trajectories_per_second": 111.21046409414787, + "compile_seconds": 2.1881700130179524 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001728332252241671, + "trajectories_per_second": 578.59245449073, + "compile_seconds": 0.7531253910856321 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009198783803731203, + "trajectories_per_second": 108.71002312222848, + "compile_seconds": 2.458911601919681 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0027445124681381616, + "trajectories_per_second": 364.36343853755045, + "compile_seconds": 0.4927656380459666 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.017914881464093924, + "trajectories_per_second": 55.819515300966955, + "compile_seconds": 2.2372332849772647 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0033871910750998984, + "trajectories_per_second": 295.2298756781848, + "compile_seconds": 0.7282182819908485 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01836333196843043, + "trajectories_per_second": 54.456348211706, + "compile_seconds": 2.5651017769705504 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005403019772428606, + "trajectories_per_second": 185.0816843393689, + "compile_seconds": 0.45458097406663 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.035812951042316854, + "trajectories_per_second": 27.92285949343835, + "compile_seconds": 2.2190580539172515 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006687036000325212, + "trajectories_per_second": 149.5430860475952, + "compile_seconds": 0.7877713339403272 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03657345997635275, + "trajectories_per_second": 27.342231242178578, + "compile_seconds": 2.669947710004635 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0012014583247946576, + "trajectories_per_second": 3329.2873480931135, + "compile_seconds": 0.7411288149887696 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007314832667664935, + "trajectories_per_second": 546.8341084112445, + "compile_seconds": 2.7222610749304295 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013529325757796566, + "trajectories_per_second": 2956.5405339544827, + "compile_seconds": 0.7850496270693839 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00741405082711329, + "trajectories_per_second": 539.5161286690864, + "compile_seconds": 2.830660925945267 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0022959496681800202, + "trajectories_per_second": 1742.1984703918906, + "compile_seconds": 0.5999633140163496 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.014622149989008904, + "trajectories_per_second": 273.5575823669363, + "compile_seconds": 2.507592984009534 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0026088268884147205, + "trajectories_per_second": 1533.2561994677383, + "compile_seconds": 0.8342320459196344 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.014867656670200327, + "trajectories_per_second": 269.04037998249686, + "compile_seconds": 3.0378022770164534 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004471874361942438, + "trajectories_per_second": 894.4795126718474, + "compile_seconds": 0.6013494469225407 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03060528391506523, + "trajectories_per_second": 130.69638599336858, + "compile_seconds": 2.5861572610447183 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005181894327203433, + "trajectories_per_second": 771.9184814327779, + "compile_seconds": 0.9482657789485529 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.030956086004152894, + "trajectories_per_second": 129.215301942997, + "compile_seconds": 2.820216372027062 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008888749778270722, + "trajectories_per_second": 450.00704258526076, + "compile_seconds": 0.6464779659872875 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.06418487604241818, + "trajectories_per_second": 62.319977020077125, + "compile_seconds": 2.5713397080544382 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.010260737006319687, + "trajectories_per_second": 389.8355447114921, + "compile_seconds": 0.9137451130663976 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.06488121603615582, + "trajectories_per_second": 61.65112561655677, + "compile_seconds": 3.0693628640146926 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009877368994057178, + "trajectories_per_second": 16198.645620738242, + "compile_seconds": 0.46920299099292606 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007081081663879256, + "trajectories_per_second": 2259.541798764493, + "compile_seconds": 2.8289907249854878 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011576078804431571, + "trajectories_per_second": 13821.605977556801, + "compile_seconds": 0.657542129047215 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007209093833807856, + "trajectories_per_second": 2219.419023923118, + "compile_seconds": 3.0945253369864076 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0019127228297293186, + "trajectories_per_second": 8365.03844222127, + "compile_seconds": 0.5078434629831463 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01422940733997772, + "trajectories_per_second": 1124.4319329483087, + "compile_seconds": 2.853143715998158 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002240234276872467, + "trajectories_per_second": 7142.11016462849, + "compile_seconds": 0.6932006339775398 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.014579707329782346, + "trajectories_per_second": 1097.4157188544098, + "compile_seconds": 3.2591083210427314 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0037042140029370785, + "trajectories_per_second": 4319.4048689718165, + "compile_seconds": 0.4679751800140366 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.028543114895001054, + "trajectories_per_second": 560.5554985451917, + "compile_seconds": 2.8852242149878293 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004348183295223862, + "trajectories_per_second": 3679.697683760191, + "compile_seconds": 0.8160997150698677 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.02922388003207743, + "trajectories_per_second": 547.4974569577239, + "compile_seconds": 3.5154968770220876 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007285461489421626, + "trajectories_per_second": 2196.1546325145973, + "compile_seconds": 0.47133537696208805 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.05705096002202481, + "trajectories_per_second": 280.4510212242374, + "compile_seconds": 2.8587830069009215 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008682303386740387, + "trajectories_per_second": 1842.8289461106829, + "compile_seconds": 0.8617396949557588 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.05849828303325921, + "trajectories_per_second": 273.51230105169407, + "compile_seconds": 4.178306789021008 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009410176859876397, + "trajectories_per_second": 34005.73706158836, + "compile_seconds": 0.4545221939915791 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006921677135064134, + "trajectories_per_second": 4623.15698573876, + "compile_seconds": 2.5349392690695822 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0010999878868460655, + "trajectories_per_second": 29091.229442309435, + "compile_seconds": 0.6856201749760658 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0070723916675585015, + "trajectories_per_second": 4524.636290547366, + "compile_seconds": 2.8466071450384334 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001806339223144783, + "trajectories_per_second": 17715.38789059175, + "compile_seconds": 0.4752070080721751 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0138525446333612, + "trajectories_per_second": 2310.0448940575243, + "compile_seconds": 2.5599693639669567 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021056555636713038, + "trajectories_per_second": 15197.16735827705, + "compile_seconds": 0.6845111540751532 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01424610634179165, + "trajectories_per_second": 2246.227792511027, + "compile_seconds": 3.022973488084972 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0034967176974392855, + "trajectories_per_second": 9151.439369393252, + "compile_seconds": 0.46306284598540515 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0278529740171507, + "trajectories_per_second": 1148.8898808542217, + "compile_seconds": 2.5829724160721526 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004175346630456095, + "trajectories_per_second": 7664.034350246143, + "compile_seconds": 0.8565303880022839 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.028483098023571074, + "trajectories_per_second": 1123.4732954090362, + "compile_seconds": 3.2766328420257196 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006881416143317308, + "trajectories_per_second": 4650.20561662673, + "compile_seconds": 0.461951993056573 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.05579029710497707, + "trajectories_per_second": 573.5764399997302, + "compile_seconds": 2.682668666006066 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008254363667219877, + "trajectories_per_second": 3876.7373585779756, + "compile_seconds": 0.8560209930874407 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.05706574197392911, + "trajectories_per_second": 560.756749901183, + "compile_seconds": 3.9734320789575577 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-a-kernels-scale.md b/benchmarks/results/vulcan-l40s-a-kernels-scale.md new file mode 100644 index 0000000..445cf37 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-a-kernels-scale.md @@ -0,0 +1,198 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=1, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | primal | 0.445 | 2.25e+03 | 0.47 | +| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | grad | 2.342 | 427 | 1.43 | +| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | primal | 0.529 | 1.89e+03 | 0.33 | +| sra1 | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | grad | 2.471 | 405 | 1.28 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | primal | 0.822 | 1.22e+03 | 0.39 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | grad | 4.553 | 220 | 1.38 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.972 | 1.03e+03 | 0.35 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 4.817 | 208 | 1.30 | +| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | primal | 1.533 | 652 | 0.36 | +| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | grad | 9.022 | 111 | 1.39 | +| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | primal | 1.850 | 541 | 0.35 | +| sra1 | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | grad | 9.489 | 105 | 1.30 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | primal | 3.002 | 333 | 0.36 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | grad | 17.824 | 56.1 | 1.37 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 3.564 | 281 | 0.37 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 18.924 | 52.8 | 1.29 | +| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | primal | 0.618 | 6.47e+03 | 0.55 | +| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | grad | 4.215 | 949 | 1.65 | +| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | primal | 0.697 | 5.74e+03 | 0.38 | +| sra1 | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | grad | 4.353 | 919 | 1.34 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | primal | 1.172 | 3.41e+03 | 0.35 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | grad | 8.355 | 479 | 1.39 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 1.307 | 3.06e+03 | 0.35 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 8.634 | 463 | 1.38 | +| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | primal | 2.258 | 1.77e+03 | 0.36 | +| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | grad | 16.870 | 237 | 1.39 | +| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | primal | 2.555 | 1.57e+03 | 0.35 | +| sra1 | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | grad | 17.002 | 235 | 1.39 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | primal | 4.365 | 916 | 0.34 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | grad | 34.251 | 117 | 1.41 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 5.010 | 798 | 0.35 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 34.530 | 116 | 1.41 | +| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.645 | 2.48e+04 | 0.49 | +| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | grad | 3.999 | 4e+03 | 1.57 | +| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.722 | 2.22e+04 | 0.35 | +| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | grad | 4.105 | 3.9e+03 | 1.40 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 1.194 | 1.34e+04 | 0.33 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | grad | 7.826 | 2.04e+03 | 1.42 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.348 | 1.19e+04 | 0.35 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 8.141 | 1.97e+03 | 1.39 | +| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 2.282 | 7.01e+03 | 0.37 | +| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | grad | 15.496 | 1.03e+03 | 1.38 | +| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 2.587 | 6.19e+03 | 0.35 | +| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | grad | 16.209 | 987 | 1.43 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 4.547 | 3.52e+03 | 0.34 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | grad | 30.837 | 519 | 1.43 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 5.058 | 3.16e+03 | 0.35 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 31.876 | 502 | 1.45 | +| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | primal | 0.650 | 4.92e+04 | 0.61 | +| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | grad | 3.941 | 8.12e+03 | 1.64 | +| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | primal | 0.724 | 4.42e+04 | 0.41 | +| sra1 | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | grad | 4.050 | 7.9e+03 | 1.51 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | primal | 1.219 | 2.62e+04 | 0.43 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | grad | 7.831 | 4.09e+03 | 1.48 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.351 | 2.37e+04 | 0.41 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 8.032 | 3.98e+03 | 1.55 | +| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | primal | 2.388 | 1.34e+04 | 0.40 | +| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | grad | 15.544 | 2.06e+03 | 1.53 | +| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | primal | 2.630 | 1.22e+04 | 0.41 | +| sra1 | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | grad | 15.981 | 2e+03 | 1.51 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | primal | 4.567 | 7.01e+03 | 0.43 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | grad | 31.250 | 1.02e+03 | 1.55 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 5.136 | 6.23e+03 | 0.42 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 32.036 | 999 | 1.55 | +| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | primal | 0.307 | 3.26e+03 | 0.18 | +| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | t1 | grad | 1.375 | 727 | 0.74 | +| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | primal | 0.392 | 2.55e+03 | 0.22 | +| em | stochastic-growth | 1 | 2 | 31 | float32 | explicit | steps | grad | 1.477 | 677 | 0.79 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | primal | 0.537 | 1.86e+03 | 0.19 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | t1 | grad | 2.618 | 382 | 0.74 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.697 | 1.43e+03 | 0.22 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 2.867 | 349 | 0.75 | +| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | primal | 1.014 | 986 | 0.19 | +| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | t1 | grad | 5.112 | 196 | 0.75 | +| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | primal | 1.303 | 768 | 0.22 | +| em | stochastic-growth | 1 | 2 | 127 | float32 | explicit | steps | grad | 5.601 | 179 | 0.76 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | primal | 1.939 | 516 | 0.19 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | t1 | grad | 10.114 | 98.9 | 0.75 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.522 | 397 | 0.23 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 10.972 | 91.1 | 0.80 | +| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | primal | 0.399 | 1e+04 | 0.23 | +| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | t1 | grad | 2.401 | 1.67e+03 | 0.92 | +| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | primal | 0.476 | 8.41e+03 | 0.24 | +| em | stochastic-growth | 4 | 2 | 31 | float32 | explicit | steps | grad | 2.551 | 1.57e+03 | 0.92 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | primal | 0.708 | 5.65e+03 | 0.22 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | t1 | grad | 4.674 | 856 | 0.92 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 0.854 | 4.68e+03 | 0.24 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 5.007 | 799 | 0.93 | +| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | primal | 1.333 | 3e+03 | 0.22 | +| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | t1 | grad | 9.231 | 433 | 0.93 | +| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | primal | 1.636 | 2.44e+03 | 0.24 | +| em | stochastic-growth | 4 | 2 | 127 | float32 | explicit | steps | grad | 9.984 | 401 | 0.96 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | primal | 2.596 | 1.54e+03 | 0.22 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | t1 | grad | 18.335 | 218 | 0.94 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 3.162 | 1.27e+03 | 0.25 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 19.791 | 202 | 0.96 | +| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.395 | 4.05e+04 | 0.22 | +| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | grad | 2.194 | 7.29e+03 | 0.90 | +| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.483 | 3.32e+04 | 0.24 | +| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | grad | 2.308 | 6.93e+03 | 0.94 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 0.710 | 2.25e+04 | 0.21 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | grad | 4.314 | 3.71e+03 | 0.94 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.875 | 1.83e+04 | 0.24 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 4.538 | 3.53e+03 | 0.93 | +| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 1.339 | 1.19e+04 | 0.22 | +| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | grad | 8.471 | 1.89e+03 | 0.91 | +| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 1.627 | 9.83e+03 | 0.27 | +| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | grad | 8.919 | 1.79e+03 | 0.92 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 2.557 | 6.26e+03 | 0.22 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | grad | 16.777 | 954 | 0.93 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.161 | 5.06e+03 | 0.24 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 17.751 | 901 | 0.95 | +| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | primal | 0.381 | 8.4e+04 | 0.24 | +| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | t1 | grad | 2.288 | 1.4e+04 | 0.94 | +| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | primal | 0.465 | 6.88e+04 | 0.26 | +| em | stochastic-growth | 32 | 2 | 31 | float32 | explicit | steps | grad | 2.402 | 1.33e+04 | 0.94 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | primal | 0.688 | 4.65e+04 | 0.23 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | t1 | grad | 4.525 | 7.07e+03 | 0.92 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.843 | 3.8e+04 | 0.25 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 4.738 | 6.75e+03 | 0.97 | +| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | primal | 1.293 | 2.47e+04 | 0.23 | +| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | t1 | grad | 8.888 | 3.6e+03 | 0.92 | +| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | primal | 1.609 | 1.99e+04 | 0.26 | +| em | stochastic-growth | 32 | 2 | 127 | float32 | explicit | steps | grad | 9.384 | 3.41e+03 | 0.94 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | primal | 2.494 | 1.28e+04 | 0.24 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | t1 | grad | 17.783 | 1.8e+03 | 0.95 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 3.072 | 1.04e+04 | 0.26 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 18.721 | 1.71e+03 | 0.95 | +| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | t1 | primal | 0.754 | 1.33e+03 | 0.54 | +| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | t1 | grad | 4.507 | 222 | 2.26 | +| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | steps | primal | 0.908 | 1.1e+03 | 0.72 | +| rk4 | neoclassical | 1 | 2 | 31 | float32 | none | steps | grad | 4.635 | 216 | 2.51 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | t1 | primal | 1.425 | 702 | 0.46 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | t1 | grad | 8.992 | 111 | 2.19 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.728 | 579 | 0.75 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 9.199 | 109 | 2.46 | +| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | t1 | primal | 2.745 | 364 | 0.49 | +| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | t1 | grad | 17.915 | 55.8 | 2.24 | +| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | steps | primal | 3.387 | 295 | 0.73 | +| rk4 | neoclassical | 1 | 2 | 127 | float32 | none | steps | grad | 18.363 | 54.5 | 2.57 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | t1 | primal | 5.403 | 185 | 0.45 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | t1 | grad | 35.813 | 27.9 | 2.22 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 6.687 | 150 | 0.79 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 36.573 | 27.3 | 2.67 | +| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | t1 | primal | 1.201 | 3.33e+03 | 0.74 | +| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | t1 | grad | 7.315 | 547 | 2.72 | +| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | steps | primal | 1.353 | 2.96e+03 | 0.79 | +| rk4 | neoclassical | 4 | 2 | 31 | float32 | none | steps | grad | 7.414 | 540 | 2.83 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | t1 | primal | 2.296 | 1.74e+03 | 0.60 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | t1 | grad | 14.622 | 274 | 2.51 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | primal | 2.609 | 1.53e+03 | 0.83 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | grad | 14.868 | 269 | 3.04 | +| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | t1 | primal | 4.472 | 894 | 0.60 | +| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | t1 | grad | 30.605 | 131 | 2.59 | +| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | steps | primal | 5.182 | 772 | 0.95 | +| rk4 | neoclassical | 4 | 2 | 127 | float32 | none | steps | grad | 30.956 | 129 | 2.82 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | t1 | primal | 8.889 | 450 | 0.65 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | t1 | grad | 64.185 | 62.3 | 2.57 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | primal | 10.261 | 390 | 0.91 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | grad | 64.881 | 61.7 | 3.07 | +| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | t1 | primal | 0.988 | 1.62e+04 | 0.47 | +| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | t1 | grad | 7.081 | 2.26e+03 | 2.83 | +| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | steps | primal | 1.158 | 1.38e+04 | 0.66 | +| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | steps | grad | 7.209 | 2.22e+03 | 3.09 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | t1 | primal | 1.913 | 8.37e+03 | 0.51 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | t1 | grad | 14.229 | 1.12e+03 | 2.85 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 2.240 | 7.14e+03 | 0.69 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | grad | 14.580 | 1.1e+03 | 3.26 | +| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | t1 | primal | 3.704 | 4.32e+03 | 0.47 | +| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | t1 | grad | 28.543 | 561 | 2.89 | +| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | steps | primal | 4.348 | 3.68e+03 | 0.82 | +| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | steps | grad | 29.224 | 547 | 3.52 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | t1 | primal | 7.285 | 2.2e+03 | 0.47 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | t1 | grad | 57.051 | 280 | 2.86 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 8.682 | 1.84e+03 | 0.86 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | grad | 58.498 | 274 | 4.18 | +| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | t1 | primal | 0.941 | 3.4e+04 | 0.45 | +| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | t1 | grad | 6.922 | 4.62e+03 | 2.53 | +| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | steps | primal | 1.100 | 2.91e+04 | 0.69 | +| rk4 | neoclassical | 32 | 2 | 31 | float32 | none | steps | grad | 7.072 | 4.52e+03 | 2.85 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | t1 | primal | 1.806 | 1.77e+04 | 0.48 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | t1 | grad | 13.853 | 2.31e+03 | 2.56 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 2.106 | 1.52e+04 | 0.68 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 14.246 | 2.25e+03 | 3.02 | +| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | t1 | primal | 3.497 | 9.15e+03 | 0.46 | +| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | t1 | grad | 27.853 | 1.15e+03 | 2.58 | +| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | steps | primal | 4.175 | 7.66e+03 | 0.86 | +| rk4 | neoclassical | 32 | 2 | 127 | float32 | none | steps | grad | 28.483 | 1.12e+03 | 3.28 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | t1 | primal | 6.881 | 4.65e+03 | 0.46 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | t1 | grad | 55.790 | 574 | 2.68 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 8.254 | 3.88e+03 | 0.86 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 57.066 | 561 | 3.97 | diff --git a/benchmarks/results/vulcan-l40s-b-primal-scale.json b/benchmarks/results/vulcan-l40s-b-primal-scale.json new file mode 100644 index 0000000..fb6d8c7 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-b-primal-scale.json @@ -0,0 +1,1164 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 1, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006354805152637489, + "trajectories_per_second": 25177.7979272258, + "compile_seconds": 0.5559755880385637 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007129469992254268, + "trajectories_per_second": 22442.06093494049, + "compile_seconds": 0.39084272703621536 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011766815370675631, + "trajectories_per_second": 13597.561868671784, + "compile_seconds": 0.35566729691345245 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013372436967781847, + "trajectories_per_second": 11964.909641039048, + "compile_seconds": 0.36642879201099277 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0022774502518586814, + "trajectories_per_second": 7025.400439347476, + "compile_seconds": 0.3300260939868167 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0026012061086172857, + "trajectories_per_second": 6150.992782538507, + "compile_seconds": 0.4029288350138813 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004493608636866239, + "trajectories_per_second": 3560.6127041713426, + "compile_seconds": 0.3373294989578426 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005107863663902713, + "trajectories_per_second": 3132.42503183318, + "compile_seconds": 0.4059816070366651 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006648337464809181, + "trajectories_per_second": 385058.67272089154, + "compile_seconds": 1.4882204619934782 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007188769645596805, + "trajectories_per_second": 356111.0073360086, + "compile_seconds": 0.4338768139714375 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011983258737018332, + "trajectories_per_second": 213631.37158106442, + "compile_seconds": 0.410205323016271 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013713163461943623, + "trajectories_per_second": 186681.94301806716, + "compile_seconds": 0.46616124699357897 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0023338496685028076, + "trajectories_per_second": 109690.01279513734, + "compile_seconds": 0.391712898039259 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002635435725096613, + "trajectories_per_second": 97137.63745485207, + "compile_seconds": 0.4829016280127689 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004540593200363219, + "trajectories_per_second": 56380.29849921847, + "compile_seconds": 0.39958338206633925 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005143783119921055, + "trajectories_per_second": 49768.81684776962, + "compile_seconds": 0.48374091600999236 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007221257656055968, + "trajectories_per_second": 1418035.540029848, + "compile_seconds": 1.3409415270434693 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008058236046465781, + "trajectories_per_second": 1270749.5711162626, + "compile_seconds": 0.47849129198584706 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001375524909235537, + "trajectories_per_second": 744443.079965087, + "compile_seconds": 0.3930518280249089 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0015292710985147183, + "trajectories_per_second": 669600.0473654048, + "compile_seconds": 0.4506659930339083 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0026392538889518213, + "trajectories_per_second": 387988.4403264747, + "compile_seconds": 0.39985329599585384 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0029785154983983375, + "trajectories_per_second": 343795.42444907344, + "compile_seconds": 0.4705504380399361 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005217596215920316, + "trajectories_per_second": 196258.95865139877, + "compile_seconds": 0.4048258389811963 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005795780627522618, + "trajectories_per_second": 176680.25513893622, + "compile_seconds": 0.46957040403503925 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0003995470896673699, + "trajectories_per_second": 40045.34237333648, + "compile_seconds": 0.2224262390518561 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00047845955766541394, + "trajectories_per_second": 33440.65291133504, + "compile_seconds": 0.24956504988949746 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007118711664313168, + "trajectories_per_second": 22475.97705102967, + "compile_seconds": 0.2086654029553756 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008631101975749646, + "trajectories_per_second": 18537.609733906935, + "compile_seconds": 0.2661267350194976 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013300093808485305, + "trajectories_per_second": 12029.99033720513, + "compile_seconds": 0.21264974900987 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00162748118080864, + "trajectories_per_second": 9831.14286584263, + "compile_seconds": 0.2523548499448225 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025666021133979987, + "trajectories_per_second": 6233.92302082115, + "compile_seconds": 0.21147427102550864 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003156211933431526, + "trajectories_per_second": 5069.368070795022, + "compile_seconds": 0.26409248902928084 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0004045593628563501, + "trajectories_per_second": 632787.2334792553, + "compile_seconds": 0.24785651196725667 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00048128006388791476, + "trajectories_per_second": 531914.8230075447, + "compile_seconds": 0.29310431494377553 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007103458220667359, + "trajectories_per_second": 360387.84497271135, + "compile_seconds": 0.25716804608237 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008685059104622765, + "trajectories_per_second": 294759.0763818059, + "compile_seconds": 0.29665725002996624 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001335371148717754, + "trajectories_per_second": 191707.00239091998, + "compile_seconds": 0.2603799499338493 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001629663723649393, + "trajectories_per_second": 157087.622608869, + "compile_seconds": 0.3043096459005028 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025974840575240946, + "trajectories_per_second": 98556.90904375273, + "compile_seconds": 0.25925671693403274 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 256, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0031994502060115337, + "trajectories_per_second": 80013.74721162862, + "compile_seconds": 0.31164434493985027 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0004323138322280807, + "trajectories_per_second": 2368649.6328892773, + "compile_seconds": 0.2987338110106066 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000522802716020394, + "trajectories_per_second": 1958673.8335920484, + "compile_seconds": 0.2993554960703477 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007981986331287771, + "trajectories_per_second": 1282888.6914853854, + "compile_seconds": 0.2597087600734085 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009497540397569537, + "trajectories_per_second": 1078173.8820105952, + "compile_seconds": 0.30768832308240235 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0015057500313560013, + "trajectories_per_second": 680059.7567166166, + "compile_seconds": 0.26700164598878473 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0018168388400226832, + "trajectories_per_second": 563616.308415784, + "compile_seconds": 0.31870555493514985 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002913577730456988, + "trajectories_per_second": 351457.93067253707, + "compile_seconds": 0.26981633191462606 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1024, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0035207956903972304, + "trajectories_per_second": 290843.34623360896, + "compile_seconds": 0.311832238920033 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009868918520320828, + "trajectories_per_second": 16212.516059439364, + "compile_seconds": 0.463355409912765 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011519278826502462, + "trajectories_per_second": 13889.758413685344, + "compile_seconds": 0.7675077740568668 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0019186041657424842, + "trajectories_per_second": 8339.396049319079, + "compile_seconds": 0.46876115596387535 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002235232497861778, + "trajectories_per_second": 7158.092062148161, + "compile_seconds": 0.7880562850041315 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0037642026929041515, + "trajectories_per_second": 4250.56812964971, + "compile_seconds": 0.47365841490682214 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004370011537420479, + "trajectories_per_second": 3661.3175647230546, + "compile_seconds": 0.9838717890670523 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007313157839234918, + "trajectories_per_second": 2187.8373681695175, + "compile_seconds": 0.46887913800310344 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008711993182078003, + "trajectories_per_second": 1836.5487283569757, + "compile_seconds": 1.005123715964146 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009848584997512045, + "trajectories_per_second": 259935.8182568063, + "compile_seconds": 0.4887423829641193 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011351239775459086, + "trajectories_per_second": 225526.02628786105, + "compile_seconds": 0.7136238029925153 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0018564381887419866, + "trajectories_per_second": 137898.47760752979, + "compile_seconds": 0.4788528170902282 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021981775228466305, + "trajectories_per_second": 116460.11176953584, + "compile_seconds": 0.7647949910024181 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0036615609961490217, + "trajectories_per_second": 69915.53609764884, + "compile_seconds": 0.5224804749013856 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004349283916367726, + "trajectories_per_second": 58860.26410844123, + "compile_seconds": 0.8992384019074962 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007238235344023754, + "trajectories_per_second": 35367.73644854837, + "compile_seconds": 0.4920805470319465 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 256, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008535689604468644, + "trajectories_per_second": 29991.718521017643, + "compile_seconds": 0.9124668670119718 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011423759502716935, + "trajectories_per_second": 896377.413894664, + "compile_seconds": 0.4886601109756157 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 31, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0012931094592990907, + "trajectories_per_second": 791889.6522147807, + "compile_seconds": 0.6905786630231887 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002228784681805833, + "trajectories_per_second": 459443.21511144016, + "compile_seconds": 0.4846900930861011 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0025102954500147868, + "trajectories_per_second": 407920.1115525936, + "compile_seconds": 0.766396596096456 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004280094639398158, + "trajectories_per_second": 239247.04621577924, + "compile_seconds": 0.5410391429904848 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 127, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004943655990064144, + "trajectories_per_second": 207134.1537635416, + "compile_seconds": 0.8720746849430725 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "t1", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008362751407548785, + "trajectories_per_second": 122447.73879989646, + "compile_seconds": 0.4981267689727247 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1024, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009781642816960812, + "trajectories_per_second": 104685.8916402511, + "compile_seconds": 0.9231459660222754 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-b-primal-scale.md b/benchmarks/results/vulcan-l40s-b-primal-scale.md new file mode 100644 index 0000000..2b9b872 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-b-primal-scale.md @@ -0,0 +1,78 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=1, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.635 | 2.52e+04 | 0.56 | +| sra1 | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.713 | 2.24e+04 | 0.39 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 1.177 | 1.36e+04 | 0.36 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.337 | 1.2e+04 | 0.37 | +| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 2.277 | 7.03e+03 | 0.33 | +| sra1 | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 2.601 | 6.15e+03 | 0.40 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 4.494 | 3.56e+03 | 0.34 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 5.108 | 3.13e+03 | 0.41 | +| sra1 | stochastic-growth | 256 | 2 | 31 | float32 | explicit | t1 | primal | 0.665 | 3.85e+05 | 1.49 | +| sra1 | stochastic-growth | 256 | 2 | 31 | float32 | explicit | steps | primal | 0.719 | 3.56e+05 | 0.43 | +| sra1 | stochastic-growth | 256 | 2 | 63 | float32 | explicit | t1 | primal | 1.198 | 2.14e+05 | 0.41 | +| sra1 | stochastic-growth | 256 | 2 | 63 | float32 | explicit | steps | primal | 1.371 | 1.87e+05 | 0.47 | +| sra1 | stochastic-growth | 256 | 2 | 127 | float32 | explicit | t1 | primal | 2.334 | 1.1e+05 | 0.39 | +| sra1 | stochastic-growth | 256 | 2 | 127 | float32 | explicit | steps | primal | 2.635 | 9.71e+04 | 0.48 | +| sra1 | stochastic-growth | 256 | 2 | 255 | float32 | explicit | t1 | primal | 4.541 | 5.64e+04 | 0.40 | +| sra1 | stochastic-growth | 256 | 2 | 255 | float32 | explicit | steps | primal | 5.144 | 4.98e+04 | 0.48 | +| sra1 | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | t1 | primal | 0.722 | 1.42e+06 | 1.34 | +| sra1 | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | steps | primal | 0.806 | 1.27e+06 | 0.48 | +| sra1 | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | t1 | primal | 1.376 | 7.44e+05 | 0.39 | +| sra1 | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | steps | primal | 1.529 | 6.7e+05 | 0.45 | +| sra1 | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | t1 | primal | 2.639 | 3.88e+05 | 0.40 | +| sra1 | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | steps | primal | 2.979 | 3.44e+05 | 0.47 | +| sra1 | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | t1 | primal | 5.218 | 1.96e+05 | 0.40 | +| sra1 | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | steps | primal | 5.796 | 1.77e+05 | 0.47 | +| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | t1 | primal | 0.400 | 4e+04 | 0.22 | +| em | stochastic-growth | 16 | 2 | 31 | float32 | explicit | steps | primal | 0.478 | 3.34e+04 | 0.25 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | t1 | primal | 0.712 | 2.25e+04 | 0.21 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.863 | 1.85e+04 | 0.27 | +| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | t1 | primal | 1.330 | 1.2e+04 | 0.21 | +| em | stochastic-growth | 16 | 2 | 127 | float32 | explicit | steps | primal | 1.627 | 9.83e+03 | 0.25 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | t1 | primal | 2.567 | 6.23e+03 | 0.21 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.156 | 5.07e+03 | 0.26 | +| em | stochastic-growth | 256 | 2 | 31 | float32 | explicit | t1 | primal | 0.405 | 6.33e+05 | 0.25 | +| em | stochastic-growth | 256 | 2 | 31 | float32 | explicit | steps | primal | 0.481 | 5.32e+05 | 0.29 | +| em | stochastic-growth | 256 | 2 | 63 | float32 | explicit | t1 | primal | 0.710 | 3.6e+05 | 0.26 | +| em | stochastic-growth | 256 | 2 | 63 | float32 | explicit | steps | primal | 0.869 | 2.95e+05 | 0.30 | +| em | stochastic-growth | 256 | 2 | 127 | float32 | explicit | t1 | primal | 1.335 | 1.92e+05 | 0.26 | +| em | stochastic-growth | 256 | 2 | 127 | float32 | explicit | steps | primal | 1.630 | 1.57e+05 | 0.30 | +| em | stochastic-growth | 256 | 2 | 255 | float32 | explicit | t1 | primal | 2.597 | 9.86e+04 | 0.26 | +| em | stochastic-growth | 256 | 2 | 255 | float32 | explicit | steps | primal | 3.199 | 8e+04 | 0.31 | +| em | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | t1 | primal | 0.432 | 2.37e+06 | 0.30 | +| em | stochastic-growth | 1024 | 2 | 31 | float32 | explicit | steps | primal | 0.523 | 1.96e+06 | 0.30 | +| em | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | t1 | primal | 0.798 | 1.28e+06 | 0.26 | +| em | stochastic-growth | 1024 | 2 | 63 | float32 | explicit | steps | primal | 0.950 | 1.08e+06 | 0.31 | +| em | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | t1 | primal | 1.506 | 6.8e+05 | 0.27 | +| em | stochastic-growth | 1024 | 2 | 127 | float32 | explicit | steps | primal | 1.817 | 5.64e+05 | 0.32 | +| em | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | t1 | primal | 2.914 | 3.51e+05 | 0.27 | +| em | stochastic-growth | 1024 | 2 | 255 | float32 | explicit | steps | primal | 3.521 | 2.91e+05 | 0.31 | +| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | t1 | primal | 0.987 | 1.62e+04 | 0.46 | +| rk4 | neoclassical | 16 | 2 | 31 | float32 | none | steps | primal | 1.152 | 1.39e+04 | 0.77 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | t1 | primal | 1.919 | 8.34e+03 | 0.47 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 2.235 | 7.16e+03 | 0.79 | +| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | t1 | primal | 3.764 | 4.25e+03 | 0.47 | +| rk4 | neoclassical | 16 | 2 | 127 | float32 | none | steps | primal | 4.370 | 3.66e+03 | 0.98 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | t1 | primal | 7.313 | 2.19e+03 | 0.47 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 8.712 | 1.84e+03 | 1.01 | +| rk4 | neoclassical | 256 | 2 | 31 | float32 | none | t1 | primal | 0.985 | 2.6e+05 | 0.49 | +| rk4 | neoclassical | 256 | 2 | 31 | float32 | none | steps | primal | 1.135 | 2.26e+05 | 0.71 | +| rk4 | neoclassical | 256 | 2 | 63 | float32 | none | t1 | primal | 1.856 | 1.38e+05 | 0.48 | +| rk4 | neoclassical | 256 | 2 | 63 | float32 | none | steps | primal | 2.198 | 1.16e+05 | 0.76 | +| rk4 | neoclassical | 256 | 2 | 127 | float32 | none | t1 | primal | 3.662 | 6.99e+04 | 0.52 | +| rk4 | neoclassical | 256 | 2 | 127 | float32 | none | steps | primal | 4.349 | 5.89e+04 | 0.90 | +| rk4 | neoclassical | 256 | 2 | 255 | float32 | none | t1 | primal | 7.238 | 3.54e+04 | 0.49 | +| rk4 | neoclassical | 256 | 2 | 255 | float32 | none | steps | primal | 8.536 | 3e+04 | 0.91 | +| rk4 | neoclassical | 1024 | 2 | 31 | float32 | none | t1 | primal | 1.142 | 8.96e+05 | 0.49 | +| rk4 | neoclassical | 1024 | 2 | 31 | float32 | none | steps | primal | 1.293 | 7.92e+05 | 0.69 | +| rk4 | neoclassical | 1024 | 2 | 63 | float32 | none | t1 | primal | 2.229 | 4.59e+05 | 0.48 | +| rk4 | neoclassical | 1024 | 2 | 63 | float32 | none | steps | primal | 2.510 | 4.08e+05 | 0.77 | +| rk4 | neoclassical | 1024 | 2 | 127 | float32 | none | t1 | primal | 4.280 | 2.39e+05 | 0.54 | +| rk4 | neoclassical | 1024 | 2 | 127 | float32 | none | steps | primal | 4.944 | 2.07e+05 | 0.87 | +| rk4 | neoclassical | 1024 | 2 | 255 | float32 | none | t1 | primal | 8.363 | 1.22e+05 | 0.50 | +| rk4 | neoclassical | 1024 | 2 | 255 | float32 | none | steps | primal | 9.782 | 1.05e+05 | 0.92 | diff --git a/benchmarks/results/vulcan-l40s-c-unroll4.json b/benchmarks/results/vulcan-l40s-c-unroll4.json new file mode 100644 index 0000000..03189f6 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-c-unroll4.json @@ -0,0 +1,780 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 4, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006971989419338677, + "trajectories_per_second": 1434.3108399250184, + "compile_seconds": 1.2450868539744988 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002124298398848623, + "trajectories_per_second": 470.7436584907296, + "compile_seconds": 7.255458320956677 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002563808891481083, + "trajectories_per_second": 390.0446727221979, + "compile_seconds": 1.1820683469995856 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008215861841260145, + "trajectories_per_second": 121.71577605869533, + "compile_seconds": 7.17904240405187 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0010690772451642083, + "trajectories_per_second": 3741.5444188839765, + "compile_seconds": 1.2497892199316993 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004303974004208364, + "trajectories_per_second": 929.3736430770393, + "compile_seconds": 6.750535171944648 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003933795417348544, + "trajectories_per_second": 1016.8296964197693, + "compile_seconds": 0.9893559309421107 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01660529401851818, + "trajectories_per_second": 240.88703250536912, + "compile_seconds": 5.85350946104154 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001099110838177419, + "trajectories_per_second": 14557.221568783469, + "compile_seconds": 1.2533364739501849 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004105353242872904, + "trajectories_per_second": 3897.3503748494213, + "compile_seconds": 6.141521341982298 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004040575159403185, + "trajectories_per_second": 3959.8322933716413, + "compile_seconds": 1.0728455058997497 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01586910302285105, + "trajectories_per_second": 1008.2485429050691, + "compile_seconds": 5.67013100697659 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011212968648495998, + "trajectories_per_second": 28538.383547779005, + "compile_seconds": 1.765444446937181 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004042536909268661, + "trajectories_per_second": 7915.821356295088, + "compile_seconds": 6.6290037749568 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0041155470826197416, + "trajectories_per_second": 7775.393977422433, + "compile_seconds": 1.361965456046164 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.015573241321059564, + "trajectories_per_second": 2054.8066610081146, + "compile_seconds": 6.725575314951129 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00046295225038193164, + "trajectories_per_second": 2160.0499817745967, + "compile_seconds": 0.7804745350731537 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0011810132515771936, + "trajectories_per_second": 846.7305499447547, + "compile_seconds": 3.3045207059476525 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0015336008145823143, + "trajectories_per_second": 652.0601648691456, + "compile_seconds": 0.7775656069861725 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004280021551742472, + "trajectories_per_second": 233.64368331109978, + "compile_seconds": 3.313947707065381 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0005910745366355, + "trajectories_per_second": 6767.336016145616, + "compile_seconds": 0.6687185120536014 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0022954671003390104, + "trajectories_per_second": 1742.5647265470511, + "compile_seconds": 3.773560242028907 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002103084047451954, + "trajectories_per_second": 1901.9686849159943, + "compile_seconds": 0.6551523020025343 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008795526390895248, + "trajectories_per_second": 454.77664692594504, + "compile_seconds": 3.7498066320549697 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006124911045557574, + "trajectories_per_second": 26122.828365980713, + "compile_seconds": 0.6658988069975749 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021762345938689327, + "trajectories_per_second": 7352.1485436710345, + "compile_seconds": 3.589312301017344 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021402357378974557, + "trajectories_per_second": 7475.811994298453, + "compile_seconds": 0.6453863370697945 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00822971984356021, + "trajectories_per_second": 1944.1731072437499, + "compile_seconds": 3.514980165986344 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00060828894015346, + "trajectories_per_second": 52606.578695852986, + "compile_seconds": 1.0748092120047659 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021829829572445965, + "trajectories_per_second": 14658.840965204337, + "compile_seconds": 3.92439767497126 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002194125268777663, + "trajectories_per_second": 14584.399740233177, + "compile_seconds": 0.942150270100683 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008260089671239257, + "trajectories_per_second": 3874.0499526803633, + "compile_seconds": 4.008594830986112 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001262728130984071, + "trajectories_per_second": 791.9361068012943, + "compile_seconds": 2.060194884077646 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004306531363082203, + "trajectories_per_second": 232.20543766905152, + "compile_seconds": 14.47477564692963 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00470992939081043, + "trajectories_per_second": 212.31740797454538, + "compile_seconds": 2.1245737449498847 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.017145550984423608, + "trajectories_per_second": 58.3241682293255, + "compile_seconds": 14.554241691017523 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021403831418692356, + "trajectories_per_second": 1868.8242874622565, + "compile_seconds": 3.1521481269737706 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007487423504547526, + "trajectories_per_second": 534.2291640870293, + "compile_seconds": 17.594308632891625 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008161732674731562, + "trajectories_per_second": 490.0920134745236, + "compile_seconds": 3.1455108739901334 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.029878524015657604, + "trajectories_per_second": 133.87542162068755, + "compile_seconds": 17.492421875009313 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0017456265545829578, + "trajectories_per_second": 9165.76340912877, + "compile_seconds": 2.474717544973828 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007120985343741874, + "trajectories_per_second": 2246.8800633133246, + "compile_seconds": 18.362305072019808 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006705917418003082, + "trajectories_per_second": 2385.9524361343165, + "compile_seconds": 2.6125571640441194 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.02837010205257684, + "trajectories_per_second": 563.9740022911454, + "compile_seconds": 24.234453483950347 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0016487168613821268, + "trajectories_per_second": 19409.033018061244, + "compile_seconds": 2.3910328299971297 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007063096001123388, + "trajectories_per_second": 4530.591116828992, + "compile_seconds": 15.935979263973422 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006319132284261286, + "trajectories_per_second": 5063.986408339739, + "compile_seconds": 2.5525463869562373 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.029351581004448235, + "trajectories_per_second": 1090.2308804132356, + "compile_seconds": 23.500358068966307 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-c-unroll4.md b/benchmarks/results/vulcan-l40s-c-unroll4.md new file mode 100644 index 0000000..cc74030 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-c-unroll4.md @@ -0,0 +1,54 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=4, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.697 | 1.43e+03 | 1.25 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 2.124 | 471 | 7.26 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.564 | 390 | 1.18 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 8.216 | 122 | 7.18 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 1.069 | 3.74e+03 | 1.25 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 4.304 | 929 | 6.75 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 3.934 | 1.02e+03 | 0.99 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 16.605 | 241 | 5.85 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.099 | 1.46e+04 | 1.25 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 4.105 | 3.9e+03 | 6.14 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 4.041 | 3.96e+03 | 1.07 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 15.869 | 1.01e+03 | 5.67 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.121 | 2.85e+04 | 1.77 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 4.043 | 7.92e+03 | 6.63 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 4.116 | 7.78e+03 | 1.36 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 15.573 | 2.05e+03 | 6.73 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.463 | 2.16e+03 | 0.78 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.181 | 847 | 3.30 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 1.534 | 652 | 0.78 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 4.280 | 234 | 3.31 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 0.591 | 6.77e+03 | 0.67 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 2.295 | 1.74e+03 | 3.77 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 2.103 | 1.9e+03 | 0.66 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 8.796 | 455 | 3.75 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.612 | 2.61e+04 | 0.67 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 2.176 | 7.35e+03 | 3.59 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 2.140 | 7.48e+03 | 0.65 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 8.230 | 1.94e+03 | 3.51 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.608 | 5.26e+04 | 1.07 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 2.183 | 1.47e+04 | 3.92 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 2.194 | 1.46e+04 | 0.94 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 8.260 | 3.87e+03 | 4.01 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.263 | 792 | 2.06 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 4.307 | 232 | 14.47 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 4.710 | 212 | 2.12 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 17.146 | 58.3 | 14.55 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | primal | 2.140 | 1.87e+03 | 3.15 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | grad | 7.487 | 534 | 17.59 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | primal | 8.162 | 490 | 3.15 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | grad | 29.879 | 134 | 17.49 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 1.746 | 9.17e+03 | 2.47 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | grad | 7.121 | 2.25e+03 | 18.36 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 6.706 | 2.39e+03 | 2.61 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | grad | 28.370 | 564 | 24.23 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 1.649 | 1.94e+04 | 2.39 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 7.063 | 4.53e+03 | 15.94 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 6.319 | 5.06e+03 | 2.55 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 29.352 | 1.09e+03 | 23.50 | diff --git a/benchmarks/results/vulcan-l40s-c-unroll8.json b/benchmarks/results/vulcan-l40s-c-unroll8.json new file mode 100644 index 0000000..9fc5cf7 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-c-unroll8.json @@ -0,0 +1,780 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 8, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006998448846085618, + "trajectories_per_second": 1428.8880607583799, + "compile_seconds": 2.1025211340747774 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0017523762295380808, + "trajectories_per_second": 570.6537118821773, + "compile_seconds": 15.505087749916129 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002500084890542846, + "trajectories_per_second": 399.98641797433885, + "compile_seconds": 2.134946339065209 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006634242716245353, + "trajectories_per_second": 150.73310440561474, + "compile_seconds": 15.41392176807858 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0010648247113244401, + "trajectories_per_second": 3756.486825915937, + "compile_seconds": 2.01313587394543 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003708232230005356, + "trajectories_per_second": 1078.681094359137, + "compile_seconds": 13.491816036985256 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0039193480818842845, + "trajectories_per_second": 1020.5778911264603, + "compile_seconds": 1.7366818400332704 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.014399717367875079, + "trajectories_per_second": 277.7832298933703, + "compile_seconds": 12.477905547944829 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001080075476784259, + "trajectories_per_second": 14813.779540331087, + "compile_seconds": 2.114171709981747 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0035481390077620745, + "trajectories_per_second": 4509.406188708405, + "compile_seconds": 12.483813125058077 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003984976414358243, + "trajectories_per_second": 4015.0802254062287, + "compile_seconds": 1.944047044031322 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.013775509665720165, + "trajectories_per_second": 1161.4815268734046, + "compile_seconds": 11.734136439976282 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001090569411065768, + "trajectories_per_second": 29342.469791746436, + "compile_seconds": 2.712940454017371 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0034320351745312414, + "trajectories_per_second": 9323.913763317028, + "compile_seconds": 13.425745286047459 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0040747205493971705, + "trajectories_per_second": 7853.299290606371, + "compile_seconds": 2.3499814469832927 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.013294546360460421, + "trajectories_per_second": 2407.0020241662287, + "compile_seconds": 13.817451639915816 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00043423763891480054, + "trajectories_per_second": 2302.8865081780823, + "compile_seconds": 1.2170277669792995 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0010025063681456706, + "trajectories_per_second": 997.4998980302672, + "compile_seconds": 7.303474535932764 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0014253398403525352, + "trajectories_per_second": 701.5870683532328, + "compile_seconds": 1.2481635819422081 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0035663329942438463, + "trajectories_per_second": 280.40006404730735, + "compile_seconds": 7.355043246061541 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0005990132657056556, + "trajectories_per_second": 6677.6484412041855, + "compile_seconds": 0.9456670730141923 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00201770125325614, + "trajectories_per_second": 1982.4540394891724, + "compile_seconds": 6.381001956993714 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002031796752514007, + "trajectories_per_second": 1968.7008530999335, + "compile_seconds": 0.9697260030079633 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0075386849930509925, + "trajectories_per_second": 530.5965169903131, + "compile_seconds": 6.386174759012647 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000614990030457689, + "trajectories_per_second": 26016.681909611525, + "compile_seconds": 1.014958688057959 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0019072774425148964, + "trajectories_per_second": 8388.921109926583, + "compile_seconds": 6.508749642991461 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0020889290406004243, + "trajectories_per_second": 7659.427241913921, + "compile_seconds": 0.985628140042536 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007178531843237579, + "trajectories_per_second": 2228.868012206778, + "compile_seconds": 6.388680926989764 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000623197469394654, + "trajectories_per_second": 51348.09040717601, + "compile_seconds": 1.2464099749922752 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001897456720471382, + "trajectories_per_second": 16864.67978676757, + "compile_seconds": 6.798548002028838 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021914461511187257, + "trajectories_per_second": 14602.229666316058, + "compile_seconds": 1.152224877034314 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007048038998618722, + "trajectories_per_second": 4540.269996558103, + "compile_seconds": 7.116658751037903 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0012470546781140216, + "trajectories_per_second": 801.88945805676, + "compile_seconds": 3.154672044911422 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0033622876425007625, + "trajectories_per_second": 297.4165527540148, + "compile_seconds": 32.872302406001836 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004600644600577653, + "trajectories_per_second": 217.3608454507529, + "compile_seconds": 3.0276695468928665 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.013169295697783431, + "trajectories_per_second": 75.93420505914476, + "compile_seconds": 32.81373755901586 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021212612378544043, + "trajectories_per_second": 1885.6706230326854, + "compile_seconds": 5.318178678979166 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006540245576096433, + "trajectories_per_second": 611.597829693027, + "compile_seconds": 40.33293925994076 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008166779787279666, + "trajectories_per_second": 489.78913405137746, + "compile_seconds": 5.438302522990853 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 4, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.025998475961387157, + "trajectories_per_second": 153.85517235474825, + "compile_seconds": 39.74205899902154 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.001730186678469181, + "trajectories_per_second": 9247.557040582658, + "compile_seconds": 4.074979388038628 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006082673120545223, + "trajectories_per_second": 2630.4224611310087, + "compile_seconds": 41.33263551397249 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006528985420508044, + "trajectories_per_second": 2450.6104654090313, + "compile_seconds": 4.321440764004365 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.024196182028390467, + "trajectories_per_second": 661.2613502918139, + "compile_seconds": 48.491429385961965 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0016249311718575913, + "trajectories_per_second": 19693.141810688627, + "compile_seconds": 4.038868790958077 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006009442877257243, + "trajectories_per_second": 5324.952853966565, + "compile_seconds": 35.891845369013026 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006131911504780874, + "trajectories_per_second": 5218.601079785729, + "compile_seconds": 4.401859620003961 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.025074374047107995, + "trajectories_per_second": 1276.2033436958634, + "compile_seconds": 45.259774502017535 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-c-unroll8.md b/benchmarks/results/vulcan-l40s-c-unroll8.md new file mode 100644 index 0000000..2ca4806 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-c-unroll8.md @@ -0,0 +1,54 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=8, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.700 | 1.43e+03 | 2.10 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.752 | 571 | 15.51 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.500 | 400 | 2.13 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 6.634 | 151 | 15.41 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 1.065 | 3.76e+03 | 2.01 | +| sra1 | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 3.708 | 1.08e+03 | 13.49 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 3.919 | 1.02e+03 | 1.74 | +| sra1 | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 14.400 | 278 | 12.48 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.080 | 1.48e+04 | 2.11 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 3.548 | 4.51e+03 | 12.48 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.985 | 4.02e+03 | 1.94 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 13.776 | 1.16e+03 | 11.73 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.091 | 2.93e+04 | 2.71 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 3.432 | 9.32e+03 | 13.43 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 4.075 | 7.85e+03 | 2.35 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 13.295 | 2.41e+03 | 13.82 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.434 | 2.3e+03 | 1.22 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.003 | 997 | 7.30 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 1.425 | 702 | 1.25 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 3.566 | 280 | 7.36 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | primal | 0.599 | 6.68e+03 | 0.95 | +| em | stochastic-growth | 4 | 2 | 63 | float32 | explicit | steps | grad | 2.018 | 1.98e+03 | 6.38 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | primal | 2.032 | 1.97e+03 | 0.97 | +| em | stochastic-growth | 4 | 2 | 255 | float32 | explicit | steps | grad | 7.539 | 531 | 6.39 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.615 | 2.6e+04 | 1.01 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 1.907 | 8.39e+03 | 6.51 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 2.089 | 7.66e+03 | 0.99 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 7.179 | 2.23e+03 | 6.39 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.623 | 5.13e+04 | 1.25 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 1.897 | 1.69e+04 | 6.80 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 2.191 | 1.46e+04 | 1.15 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 7.048 | 4.54e+03 | 7.12 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.247 | 802 | 3.15 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 3.362 | 297 | 32.87 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 4.601 | 217 | 3.03 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 13.169 | 75.9 | 32.81 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | primal | 2.121 | 1.89e+03 | 5.32 | +| rk4 | neoclassical | 4 | 2 | 63 | float32 | none | steps | grad | 6.540 | 612 | 40.33 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | primal | 8.167 | 490 | 5.44 | +| rk4 | neoclassical | 4 | 2 | 255 | float32 | none | steps | grad | 25.998 | 154 | 39.74 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | primal | 1.730 | 9.25e+03 | 4.07 | +| rk4 | neoclassical | 16 | 2 | 63 | float32 | none | steps | grad | 6.083 | 2.63e+03 | 41.33 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | primal | 6.529 | 2.45e+03 | 4.32 | +| rk4 | neoclassical | 16 | 2 | 255 | float32 | none | steps | grad | 24.196 | 661 | 48.49 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 1.625 | 1.97e+04 | 4.04 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 6.009 | 5.32e+03 | 35.89 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 6.132 | 5.22e+03 | 4.40 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 25.074 | 1.28e+03 | 45.26 | diff --git a/benchmarks/results/vulcan-l40s-d-noisemode.json b/benchmarks/results/vulcan-l40s-d-noisemode.json new file mode 100644 index 0000000..f46ba48 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-d-noisemode.json @@ -0,0 +1,268 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 1, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013318749124129467, + "trajectories_per_second": 12013.14016119797, + "compile_seconds": 0.6122790339868516 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013370389707334754, + "trajectories_per_second": 11966.741695810624, + "compile_seconds": 0.7160951240221038 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008066055675347647, + "trajectories_per_second": 1983.621319265243, + "compile_seconds": 1.5805910240160301 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007812453336858501, + "trajectories_per_second": 2048.0122325356288, + "compile_seconds": 1.5591626500245184 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.00505731444961081, + "trajectories_per_second": 3163.7344601404593, + "compile_seconds": 0.40626175701618195 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0050121010022444856, + "trajectories_per_second": 3192.2740568945014, + "compile_seconds": 0.7446511259768158 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03158391290344298, + "trajectories_per_second": 506.58700994124865, + "compile_seconds": 1.4282209059456363 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.030595276970416307, + "trajectories_per_second": 522.9565339601595, + "compile_seconds": 1.6230945909628645 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008508051179495512, + "trajectories_per_second": 18805.716682288134, + "compile_seconds": 0.22564140998292714 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008544241986237467, + "trajectories_per_second": 18726.06139406141, + "compile_seconds": 0.3509363259654492 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004469206696376205, + "trajectories_per_second": 3580.0537068409435, + "compile_seconds": 0.9315736360149458 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004436203818344934, + "trajectories_per_second": 3606.6873063486305, + "compile_seconds": 1.0160666520241648 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0031061899305010834, + "trajectories_per_second": 5151.005044118122, + "compile_seconds": 0.22345678601413965 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003116562000165383, + "trajectories_per_second": 5133.862249219154, + "compile_seconds": 0.3596626119688153 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.017609555507078767, + "trajectories_per_second": 908.597607337008, + "compile_seconds": 0.9339326129993424 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 16, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "key", + "remat": "none", + "oom": false, + "seconds_per_call": 0.01742742903297767, + "trajectories_per_second": 918.0929653894119, + "compile_seconds": 1.0483848019503057 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-d-noisemode.md b/benchmarks/results/vulcan-l40s-d-noisemode.md new file mode 100644 index 0000000..690f1b8 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-d-noisemode.md @@ -0,0 +1,22 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=1, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 1.332 | 1.2e+04 | 0.61 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | primal | 1.337 | 1.2e+04 | 0.72 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 8.066 | 1.98e+03 | 1.58 | +| sra1 | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | grad | 7.812 | 2.05e+03 | 1.56 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 5.057 | 3.16e+03 | 0.41 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | primal | 5.012 | 3.19e+03 | 0.74 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 31.584 | 507 | 1.43 | +| sra1 | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | grad | 30.595 | 523 | 1.62 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | primal | 0.851 | 1.88e+04 | 0.23 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | primal | 0.854 | 1.87e+04 | 0.35 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | explicit | steps | grad | 4.469 | 3.58e+03 | 0.93 | +| em | stochastic-growth | 16 | 2 | 63 | float32 | key | steps | grad | 4.436 | 3.61e+03 | 1.02 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | primal | 3.106 | 5.15e+03 | 0.22 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | primal | 3.117 | 5.13e+03 | 0.36 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | explicit | steps | grad | 17.610 | 909 | 0.93 | +| em | stochastic-growth | 16 | 2 | 255 | float32 | key | steps | grad | 17.427 | 918 | 1.05 | diff --git a/benchmarks/results/vulcan-l40s-e-matmul-highest.json b/benchmarks/results/vulcan-l40s-e-matmul-highest.json new file mode 100644 index 0000000..b21f9b0 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-e-matmul-highest.json @@ -0,0 +1,140 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": "highest", + "scan_unroll": 1, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0018264910009188147, + "trajectories_per_second": 17519.933021242607, + "compile_seconds": 0.4346738619497046 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009604040393605828, + "trajectories_per_second": 3331.9310090891477, + "compile_seconds": 1.2720381129765883 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.007048667857556471, + "trajectories_per_second": 4539.864928618341, + "compile_seconds": 0.3383073400473222 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03859191294759512, + "trajectories_per_second": 829.1892667630539, + "compile_seconds": 1.31303374201525 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0028107422403991222, + "trajectories_per_second": 11384.893121845294, + "compile_seconds": 0.7347728699678555 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.016519141655104857, + "trajectories_per_second": 1937.1466549601955, + "compile_seconds": 2.9693069600034505 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.011021795507986099, + "trajectories_per_second": 2903.338206267178, + "compile_seconds": 0.8703918239334598 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.06651889497879893, + "trajectories_per_second": 481.0663197306438, + "compile_seconds": 3.8560943111078814 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-e-matmul-highest.md b/benchmarks/results/vulcan-l40s-e-matmul-highest.md new file mode 100644 index 0000000..f57d609 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-e-matmul-highest.md @@ -0,0 +1,14 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=highest, scan_unroll=1, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.826 | 1.75e+04 | 0.43 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 9.604 | 3.33e+03 | 1.27 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 7.049 | 4.54e+03 | 0.34 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 38.592 | 829 | 1.31 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 2.811 | 1.14e+04 | 0.73 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 16.519 | 1.94e+03 | 2.97 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 11.022 | 2.9e+03 | 0.87 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 66.519 | 481 | 3.86 | diff --git a/benchmarks/results/vulcan-l40s-f-cudagraph.json b/benchmarks/results/vulcan-l40s-f-cudagraph.json new file mode 100644 index 0000000..0e8940a --- /dev/null +++ b/benchmarks/results/vulcan-l40s-f-cudagraph.json @@ -0,0 +1,396 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 1, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0009622347191907465, + "trajectories_per_second": 1039.247472634343, + "compile_seconds": 0.4700410889927298 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005052514995137851, + "trajectories_per_second": 197.92123347725294, + "compile_seconds": 1.5074586200062186 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003587769236988746, + "trajectories_per_second": 278.7247266881944, + "compile_seconds": 0.3678640240104869 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.019817276508547366, + "trajectories_per_second": 50.46102069417517, + "compile_seconds": 1.3597999310586601 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0013455422789168854, + "trajectories_per_second": 23782.23300850783, + "compile_seconds": 0.6324994110036641 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008021229334796468, + "trajectories_per_second": 3989.413425842658, + "compile_seconds": 1.570612299023196 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.005146991215749747, + "trajectories_per_second": 6217.224521790573, + "compile_seconds": 0.4534137600567192 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03185791301075369, + "trajectories_per_second": 1004.4600218852487, + "compile_seconds": 1.5433854449074715 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006954732269867126, + "trajectories_per_second": 1437.869872191508, + "compile_seconds": 0.2018867180449888 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0030720798749825917, + "trajectories_per_second": 325.51236969568265, + "compile_seconds": 0.7667697980068624 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002501146420591364, + "trajectories_per_second": 399.8166567807585, + "compile_seconds": 0.2249728940660134 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.011790249263867736, + "trajectories_per_second": 84.81584889511952, + "compile_seconds": 0.7430006139911711 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0008435127194644066, + "trajectories_per_second": 37936.59450721572, + "compile_seconds": 0.2570032199146226 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004686483694240451, + "trajectories_per_second": 6828.147090179157, + "compile_seconds": 0.9227630530949682 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003072791689191945, + "trajectories_per_second": 10413.982865338676, + "compile_seconds": 0.24970263394061476 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.018492621486075222, + "trajectories_per_second": 1730.4198879588657, + "compile_seconds": 0.9388698430266231 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0017479369237732429, + "trajectories_per_second": 572.1030240847115, + "compile_seconds": 0.7169918889412656 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.009450687794014812, + "trajectories_per_second": 105.81240453560504, + "compile_seconds": 2.523357450030744 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006671004501792292, + "trajectories_per_second": 149.90246217512384, + "compile_seconds": 0.8463023368967697 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.03746189596131444, + "trajectories_per_second": 26.69379043262157, + "compile_seconds": 2.656470412039198 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021086527101163354, + "trajectories_per_second": 15175.56677136964, + "compile_seconds": 0.6879883620422333 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.014226612984202802, + "trajectories_per_second": 2249.3055821180155, + "compile_seconds": 3.116745045990683 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.008247273838302741, + "trajectories_per_second": 3880.0700240342067, + "compile_seconds": 0.8314209800446406 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.05689899192657322, + "trajectories_per_second": 562.40012197923, + "compile_seconds": 3.9393029569182545 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-f-cudagraph.md b/benchmarks/results/vulcan-l40s-f-cudagraph.md new file mode 100644 index 0000000..3f70297 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-f-cudagraph.md @@ -0,0 +1,30 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=1, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.962 | 1.04e+03 | 0.47 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 5.053 | 198 | 1.51 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 3.588 | 279 | 0.37 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 19.817 | 50.5 | 1.36 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.346 | 2.38e+04 | 0.63 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 8.021 | 3.99e+03 | 1.57 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 5.147 | 6.22e+03 | 0.45 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 31.858 | 1e+03 | 1.54 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.695 | 1.44e+03 | 0.20 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 3.072 | 326 | 0.77 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.501 | 400 | 0.22 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 11.790 | 84.8 | 0.74 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.844 | 3.79e+04 | 0.26 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 4.686 | 6.83e+03 | 0.92 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 3.073 | 1.04e+04 | 0.25 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 18.493 | 1.73e+03 | 0.94 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.748 | 572 | 0.72 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 9.451 | 106 | 2.52 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 6.671 | 150 | 0.85 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 37.462 | 26.7 | 2.66 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 2.109 | 1.52e+04 | 0.69 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 14.227 | 2.25e+03 | 3.12 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 8.247 | 3.88e+03 | 0.83 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 56.899 | 562 | 3.94 | diff --git a/benchmarks/results/vulcan-l40s-g-unroll8-promoted.json b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.json new file mode 100644 index 0000000..f2cb095 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.json @@ -0,0 +1,396 @@ +{ + "config": { + "device": "NVIDIA L40S", + "backend": "gpu", + "jax_version": "0.9.1", + "matmul_precision": null, + "scan_unroll": 8, + "remat": "none" + }, + "records": [ + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0007148562263782052, + "trajectories_per_second": 1398.8826887141574, + "compile_seconds": 2.184549228986725 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0017632524977670982, + "trajectories_per_second": 567.1337492879517, + "compile_seconds": 15.08952729194425 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.002500937333227032, + "trajectories_per_second": 399.85008289258934, + "compile_seconds": 2.118349283002317 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006678907865924495, + "trajectories_per_second": 149.72507782327074, + "compile_seconds": 15.132703572977334 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0010891472568822114, + "trajectories_per_second": 29380.783725795787, + "compile_seconds": 2.608467624988407 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003410756928912763, + "trajectories_per_second": 9382.081651359584, + "compile_seconds": 13.017872113967314 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004007596903565255, + "trajectories_per_second": 7984.8349946403105, + "compile_seconds": 2.1188945450121537 + }, + { + "solver": "sra1", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.013283782677414516, + "trajectories_per_second": 2408.9523878170153, + "compile_seconds": 13.366022419999354 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.000445107000857769, + "trajectories_per_second": 2246.650800982444, + "compile_seconds": 1.2408874860266224 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0010245928041091778, + "trajectories_per_second": 975.9974850393763, + "compile_seconds": 7.035418029059656 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0014510225695890508, + "trajectories_per_second": 689.169156261445, + "compile_seconds": 1.2501748150680214 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003635078767099633, + "trajectories_per_second": 275.09720258355856, + "compile_seconds": 7.035499937948771 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0006264614267274737, + "trajectories_per_second": 51080.559208828665, + "compile_seconds": 1.2203648679424077 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0018924760855346297, + "trajectories_per_second": 16909.06439695375, + "compile_seconds": 6.612741588032804 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0021858520446006546, + "trajectories_per_second": 14639.600186592801, + "compile_seconds": 1.1961514059221372 + }, + { + "solver": "em", + "drift": "stochastic-growth", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "explicit", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0070658908419621485, + "trajectories_per_second": 4528.799087860494, + "compile_seconds": 6.868361442000605 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0012508814995128073, + "trajectories_per_second": 799.436237876633, + "compile_seconds": 2.7070158420829102 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.003360000147949904, + "trajectories_per_second": 297.6190345140751, + "compile_seconds": 32.03250827000011 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.004642306303139776, + "trajectories_per_second": 215.41017216456837, + "compile_seconds": 3.114756237016991 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 1, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.013239547687893113, + "trajectories_per_second": 75.53128124720217, + "compile_seconds": 32.787103537935764 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0016123186422711505, + "trajectories_per_second": 19847.193452358795, + "compile_seconds": 3.9508127509616315 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 63, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.0060252813855186105, + "trajectories_per_second": 5310.955282007246, + "compile_seconds": 34.89513242407702 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "primal", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.006146438128780574, + "trajectories_per_second": 5206.267325812756, + "compile_seconds": 4.017467572935857 + }, + { + "solver": "rk4", + "drift": "neoclassical", + "batch": 32, + "dim": 2, + "n_steps": 255, + "dtype": "float32", + "save": "steps", + "mode": "grad", + "noise": "none", + "remat": "none", + "oom": false, + "seconds_per_call": 0.025152009911835194, + "trajectories_per_second": 1272.264129672695, + "compile_seconds": 44.808164507034235 + } + ] +} diff --git a/benchmarks/results/vulcan-l40s-g-unroll8-promoted.md b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.md new file mode 100644 index 0000000..d947451 --- /dev/null +++ b/benchmarks/results/vulcan-l40s-g-unroll8-promoted.md @@ -0,0 +1,30 @@ +# gpu_trajectories — NVIDIA L40S (gpu, jax 0.9.1) + +matmul_precision=None, scan_unroll=8, remat=none + +| solver | drift | B | d | n_steps | dtype | noise | save | mode | ms/call | traj/s | compile s | +|---|---|---:|---:|---:|---|---|---|---|---:|---:|---:| +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.715 | 1.4e+03 | 2.18 | +| sra1 | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.763 | 567 | 15.09 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 2.501 | 400 | 2.12 | +| sra1 | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 6.679 | 150 | 15.13 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 1.089 | 2.94e+04 | 2.61 | +| sra1 | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 3.411 | 9.38e+03 | 13.02 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 4.008 | 7.98e+03 | 2.12 | +| sra1 | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 13.284 | 2.41e+03 | 13.37 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | primal | 0.445 | 2.25e+03 | 1.24 | +| em | stochastic-growth | 1 | 2 | 63 | float32 | explicit | steps | grad | 1.025 | 976 | 7.04 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | primal | 1.451 | 689 | 1.25 | +| em | stochastic-growth | 1 | 2 | 255 | float32 | explicit | steps | grad | 3.635 | 275 | 7.04 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | primal | 0.626 | 5.11e+04 | 1.22 | +| em | stochastic-growth | 32 | 2 | 63 | float32 | explicit | steps | grad | 1.892 | 1.69e+04 | 6.61 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | primal | 2.186 | 1.46e+04 | 1.20 | +| em | stochastic-growth | 32 | 2 | 255 | float32 | explicit | steps | grad | 7.066 | 4.53e+03 | 6.87 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | primal | 1.251 | 799 | 2.71 | +| rk4 | neoclassical | 1 | 2 | 63 | float32 | none | steps | grad | 3.360 | 298 | 32.03 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | primal | 4.642 | 215 | 3.11 | +| rk4 | neoclassical | 1 | 2 | 255 | float32 | none | steps | grad | 13.240 | 75.5 | 32.79 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | primal | 1.612 | 1.98e+04 | 3.95 | +| rk4 | neoclassical | 32 | 2 | 63 | float32 | none | steps | grad | 6.025 | 5.31e+03 | 34.90 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | primal | 6.146 | 5.21e+03 | 4.02 | +| rk4 | neoclassical | 32 | 2 | 255 | float32 | none | steps | grad | 25.152 | 1.27e+03 | 44.81 | diff --git a/docs/adaptive_ad.md b/docs/adaptive_ad.md deleted file mode 100644 index 591228e..0000000 --- a/docs/adaptive_ad.md +++ /dev/null @@ -1,261 +0,0 @@ -# Adaptive Stepping and AD - -The default `adaptive_loop="bounded"` supports forward mode, reverse mode, and -reverse-over-forward through one adaptive solve. The alternative -`adaptive_loop="forward"` executes a data-dependent number of attempts and -supports primal evaluation, JVP, and nested forward mode, but not ordinary -reverse mode. Both paths use the same deliberate frozen-controller derivative -convention described below. - -## Default tolerances follow the state precision - -`IController()` and `PIController()` resolve omitted tolerances from the -state dtype at trace time: - -| State/time dtype | `rtol` | `atol` | unit-scale `dt_min` | -|---|---:|---:|---:| -| float32 | `1e-4` | `1e-6` | `10 * eps` ≈ `1.19e-6` | -| float64 | `1e-7` | `1e-9` | `10 * eps` ≈ `2.22e-15` | - -The policy never changes `jax_enable_x64`: a float32 `x_0` stays float32 even -when x64 is enabled globally. Explicit tolerances override the defaults and -are cast to the state dtype. Use explicit values whenever tolerances are part -of a reproducibility or accuracy contract. - -Automatic `dt_min` is `10 * eps * max(1, abs(t_1))`; set it explicitly when -the relevant time scale differs materially from the absolute horizon. The -controller also floors exact-zero error ratios at machine epsilon before -applying a negative power; the growth-factor clip then selects maximal step -growth without introducing infinities. - -## The controller is stop-gradiented - -The adaptive step-size controllers (`IController` and `PIController`) compute -their scaled error norms, decisions, and next-step factors inside -`stop_gradient`. Accept/reject is a discrete branch, and differentiating -`E**(-1/order)` is singular at the exact-zero error of a flat-start policy. - -States remain differentiable through solver stages, but the derivative holds -the realized step sizes and accept/reject pattern fixed. It is the derivative -of the discrete flow on a **frozen mesh**, not the total derivative of a mesh -that moves with parameters. In particular, for parameter-only differentiation, -the returned adaptive knot times in `SaveAt(steps=True)` have exactly zero -tangent. - -The `E**(-1/5)` blow-up is not hypothetical: a policy initialized flat gives -an exactly-zero error estimate on the first step, and without the -`stop_gradient` (plus the machine-epsilon error floor) the backward pass is NaN -from iteration one. `tests/test_ad.py::test_grad_finite_on_flat_field` pins -this. - -`PIController` additionally carries the previous accepted error ratio. Its -step-size factor is - -```text -safety * E_n**(-(p_coeff + i_coeff) / order) - * E_prev**(p_coeff / order) -``` - -and `E_prev` changes only after acceptance. The whole recurrence is -controller-internal and stop-gradiented. Setting `p_coeff=0, i_coeff=1` -reproduces `IController` bit for bit. - -`IController` remains the default. On the smooth ODE and DAE screening -problems, the default PI coefficients did not reduce rejections, and -matched-work accuracy was comparable or favored I, except that PI sometimes -improved common-grid ODE interpolation. PI remains useful as an opt-in for -genuinely rejection-prone or step-size-oscillatory problems; those cases were -not covered by this screen. - -## The horizon clip is the growth guard - -Every attempt is clipped so it cannot step past `t_1`, and — deliberately -unlike diffrax — the controller's next-step proposal is computed from the -**clipped** step: - -> the horizon clip doubles as the guard on step growth: without it, a -> near-flat vector field lets steps quintuple into quarter-horizon leaps -> whose Gauss–Newton linearization stalls a trust-region optimizer -> differentiating through the rollout. - -With `factor_max = 5`, a flat field would otherwise reach `dt ≈ t_1/4` within -a few accepted steps; residuals sampled from three or four giant steps make -the optimizer's linear model useless. Clipping first means the proposal can -never exceed `factor_max × remaining horizon`. - -Adaptive steps use `min(dt, remaining)` and are never enlarged to consume a -floating-point sliver. Fixed-step times are instead formed arithmetically as -`t_0 + i * dt_0`, avoiding drift from repeated addition, with a small local -endpoint snap capped below a quarter step. No time tolerance scales with -`max_steps`: changing a nonbinding attempt budget must not stretch a step or -change the numerical method. - -## Frozen mesh versus a moving-mesh derivative - -Let a parameter-dependent controller choose an internal knot -\(t_k(\theta)\), and let \(q(x,t,\theta)\) be a saved quantity. Along the exact -trajectory, its total derivative contains - -$$ -\frac{d q_k}{d\theta} -=q_x\left(x_\theta\rvert_t+f\,t_{k,\theta}\right) - +q_t t_{k,\theta}+q_\theta. -$$ - -tinydiffeq's frozen-mesh rule returns the terms holding \(t_k\) fixed. The -omitted mesh-motion term is - -$$ -\left(q_t+q_x f\right)t_{k,\theta} -=\frac{d q}{dt}\,t_{k,\theta}. -$$ - -For a collocation residual \(r(x_k,t_k,\theta)\), this omission is the total -trajectory derivative \((d r/dt)t_{k,\theta}\). If the learned policy approaches -a smooth root for which the residual vanishes along the trajectory, then -\(d r/dt\) vanishes with it locally; under that regularity, the omitted term is -\(O(\lVert r\rVert)\), so the frozen-mesh residual Jacobian is asymptotically -exact near the root. Pointwise cancellation at a few knots alone is not enough -for that conclusion. - -Far from a root the omitted term can be material, so a Gauss--Newton or LM -predicted reduction and trust ratio may be misleading. A targeted -finite-difference audit found that adaptive `SaveAt(steps=True)` state -derivatives need not approach the moving-mesh derivative under tolerance -refinement. Endpoint output at fixed `t_1` and fixed requested-grid output do -not have this particular moving-output-time ambiguity and did converge in the -same audit. - -## Interpolation knots are non-differentiable - -`SaveAt(ts=...)` brackets each query with `searchsorted` — integer indices, -no gradient. Values differentiate through the bracketing states and -derivatives (`xL`, `xR`, `fL`, `fR`) while internal knot locations remain -frozen. The requested times themselves are fixed outputs, unlike -`SaveAt(steps=True)`'s adaptive internal times. - -Zero-width brackets (duplicate rows from rejections and the frozen tail) use -the **double-where trick**: the divisor is replaced by 1 *before* dividing, - -```python -width_safe = jnp.where(degenerate, 1.0, width) -s = jnp.clip((tau - t_left) / width_safe, 0.0, 1.0) -value = jnp.where(degenerate, x_left, hermite(s, ...)) -``` - -so neither the primal nor its jvp/vjp ever evaluates `0/0`. A single `where` -on the output is not enough — reverse mode differentiates both branches, and -`NaN * 0 = NaN`. - -## What this buys you - -- With the default bounded loop, `jax.grad`, `jax.jvp`, and - `jax.grad(jax.jvp(...))` (the - Levenberg–Marquardt geodesic-acceleration pattern) all work through - adaptive solves and interpolated output, verified against closed forms in - `tests/test_ad.py`. -- With `adaptive_loop="forward"`, primal evaluation, `jax.jvp`, and nested - forward mode work; `jax.vjp`, `jax.grad`, and reverse-over-forward do not. -- `jax.vmap` over `x_0` or `p` gives genuinely per-lane adaptivity: each lane - accepts/rejects independently. Execution continues until the slowest lane - finishes. - -The dynamic loop is most useful when actual attempts are much fewer than -`max_steps`. When an integration uses most of its budget, primal performance -can be similar or slower than the bounded scan; compilation and forward-mode -paths may still improve. Choose from representative end-to-end measurements, -not from the loop form alone. - -On float32 GPUs, different loop lowerings can also fuse neural-network matrix -products differently. With accelerator-default reduced-precision matrix -multiplication, that roundoff was large enough in one width-32 policy test to -move the adaptive mesh even though both loop contracts were correct. Setting -`jax_default_matmul_precision="highest"` restored matching paths. tinydiffeq -does not change this application-global JAX setting; set it in reproducibility -entry points when adaptive decisions depend on neural-network outputs. - -## Reusing a linearization - -A single `jax.jvp` must evaluate the primal trajectory and propagate its -tangent together. If several tangent directions are needed at the same -`x_0`, `p`, and other primal inputs, cache the linearization instead: - -```python -value, pushforward = jax.linearize(endpoint, x_0) -tangent_batch = jax.jit(jax.vmap(pushforward))(directions) - -value, pullback = jax.vjp(endpoint, x_0) -cotangent_batch = jax.jit( - jax.vmap(lambda cotangent: pullback(cotangent)[0]) -)(cotangents) -``` - -The setup computes and stores residuals for that one primal point. Reuse the -pushforward or pullback only while every primal input is unchanged; otherwise -linearize again. Rodas5P's custom linear-solve rule explicitly retains and -reuses its primal pivoted LU factors. Semi-explicit DAE roots have a different -boundary: nlls-gram supplies both the primal LM root and its implicit derivative, -selecting direct `LU()` by default for a square residual system. The primal LM -Cholesky factors are not reused because they represent damped normal equations, -not the accepted root's nonsymmetric Jacobian. `jax.linearize` still avoids -repeating the primal trajectory, but the API makes no promise that the DAE -root's direct factorization is cached across pushforward directions. - -On the 256-state fixed Tsit5 benchmark, cached pushforwards were about 8–18% -faster than a fused `vmap(jvp)` on CPU after excluding setup. On the RTX 3090, -cached pushforwards and pullbacks were roughly 2.2–2.5x faster for 1–16 -directions because the primal trajectory was not replicated across mapped -lanes. For one direction at a new primal point, ordinary `jax.jvp` or -`jax.vjp` remains the right interface. - -## What to watch - -- Reductions over the valid prefix of `SaveAt(steps=True)` are - **discontinuous** in the inputs: its length changes when an accept flips to - a reject. Use `sol.accepted` when padding must not contribute. Because the - default tail repeats the endpoint, `xs[-1]` remains the reached final state. -- Adaptive `SaveAt(steps=True)` derivatives hold the internal times fixed and - omit mesh motion. Use a fixed requested grid when the residual definition - requires fixed sample times, or treat the frozen-mesh Jacobian explicitly as - a near-root approximation. -- Finite-difference checks of adaptive solves are noisy for the same reason; - compare AD against closed forms or use fixed-step solvers for FD tests. - -## Custom-rule audit - -tinydiffeq uses hand-coded derivative boundaries where they enforce the intended -mathematical derivative or avoid differentiating iteration/factorization work: - -There is no custom reverse rule that replays an adaptive solve on a recorded -mesh. The bounded path uses ordinary traced AD through its scan, subject to the -controller `stop_gradient`; the dynamic forward path deliberately exposes -JAX's no-reverse boundary. A replay/custom-VJP design would be a distinct future -API and would not restore the omitted moving-mesh term by itself. - -- Rodas5P's factored linear solve has a custom JVP - \(\delta x=A^{-1}(\delta b-\delta A\,x)\). Every stage tangent and the - transposed VJP reuse the attempt's pivoted LU factors; pivot selection is not - differentiated. -- Semi-explicit DAE roots delegate both the primal solve and implicit AD to - nlls-gram. Its implicit-function rule is - \(\delta z=-g_z^{-1}(g_y\delta y+g_t\delta t+g_p\delta p)\); the default - square `LU()` implementation differentiates the defining constraint rather - than the optimizer iterations, and reverse mode transposes that rule. - `LMRootSolver` requires residual-only convergence (`gtol=xtol=0`) and never - accepts `MAX_STEPS` as a differentiable root. `ad_solver` remains an - nlls-owned option forwarded through `solver_options`. -- Dense linear exponential actions use a Fréchet custom JVP for active matrix - or time tangents and reuse the matrix exponential when only the initial state - varies. -- Matrix-free terminal exponential sensitivities expose - `jvp_linear_ode`/`vjp_linear_ode`, applying the forward or transposed - exponential directly instead of differentiating Arnoldi orthogonalization. -- `AdaptiveKrylovExponential` uses a bounded scan and residual-controlled - internal slices. Ordinary AD follows the realized controller path; the - hand-coded initial-state rules apply independent adaptive forward or - transposed actions and avoid differentiating the Arnoldi basis. - -Explicit Runge--Kutta, Euler--Maruyama, Hermite/Rodas interpolation, and -cumulative trapezoids remain ordinary JAX programs. Their recurrences and -polynomials already transpose efficiently, and a custom rule would either -duplicate JAX's work or introduce a different continuous-adjoint derivative. diff --git a/docs/api.md b/docs/api.md index 7a2fc07..8f4a596 100644 --- a/docs/api.md +++ b/docs/api.md @@ -36,6 +36,10 @@ ::: tinydiffeq.EulerMaruyama +::: tinydiffeq.Milstein + +::: tinydiffeq.SRA1 + ::: tinydiffeq.DenseExponential ::: tinydiffeq.KrylovExponential @@ -78,6 +82,8 @@ ## Utilities +::: tinydiffeq.diagonal_brownian_increments + ::: tinydiffeq.hermite_interpolate ::: tinydiffeq.cumulative_trapezoid diff --git a/docs/aux.md b/docs/aux.md deleted file mode 100644 index a0cfc1c..0000000 --- a/docs/aux.md +++ /dev/null @@ -1,158 +0,0 @@ -# Auxiliary Outputs - -Auxiliary outputs let a model expose quantities it already computes without -adding them to the dynamical state. They are ordinary JAX pytrees, are saved -with the solution, and participate in JVP and VJP. The no-aux code path remains -available explicitly with `has_aux=False`. - -## ODE contract - -An ODE field may return only its derivative, - -```python -def f(x, t, args, p): - return dx -``` - -or pair the derivative with saved output: - -```python -def f(x, t, args, p): - return dx, {"flow": flow, "moment": moment} -``` - -The first item must have the state pytree structure and dtype. Saved aux must -be a nonempty pytree of nonempty real floating arrays; its leaves may have -different floating dtypes. The result is available as `sol.aux` with the same -leading saved-time axis as `sol.xs`. - -`has_aux=None` (the default) performs one `jax.eval_shape` trace to detect the -contract. A complete field output matching the state pytree takes precedence, -so a two-item tuple state is not misclassified as `(dx, aux)`. Set -`has_aux=True` to require aux, or `has_aux=False` to skip detection and select -the minimal value-only path. - -## DAE contracts - -For a semi-explicit DAE, saved output belongs to the differential field: - -```python -def f(y, z, t, args, p): - return dy, saved_aux - - -def g(y, z, t, args, p): - return residual -``` - -The algebraic function may additionally expose internal context that avoids -recomputing quantities already formed while evaluating the residual: - -```python -def g(y, z, t, args, p): - return residual, algebraic_aux - - -def f(y, z, t, args, p, algebraic_aux): - return dy, saved_aux -``` - -All four combinations are supported: neither aux, saved aux only, algebraic -aux only, or both. When algebraic aux is present, the differential function -must use the full six-argument signature. `has_algebraic_aux=None` detects the -form; explicit `False` avoids that trace. - -Algebraic aux is internal context. It is passed to `f`, included in the -implicit derivative path, and ignored by the nonlinear solver's residual -interface. It is not independently stored or interpolated. It may be a -nonempty pytree of bool, integer, real, or complex arrays; every inexact leaf -must be finite. This broader dtype contract is useful for cached masks and -indices that should not become solution output. - -Only the differential field's real-floating `saved_aux` becomes `sol.aux`. -This separation keeps one unambiguous output to differentiate and interpolate: -the algebraic function shares work with the dynamics, while the differential -function decides what users retain. - -## SDE and SDAE contracts - -For an SDE, only the drift owns saved aux: - -```python -def drift(x, t, args, p): - return drift_value, saved_aux - - -def diffusion(x, t, args, p): - return diffusion_value -``` - -For an SDAE, algebraic aux is passed to both stochastic fields, and the drift -still owns saved aux: - -```python -def g(y, z, t, args, p): - return residual, algebraic_aux - - -def drift(y, z, t, args, p, algebraic_aux): - return drift_value, saved_aux - - -def diffusion(y, z, t, args, p, algebraic_aux): - return diffusion_value -``` - -The fixed random key defines the path for both primal evaluation and AD. -Aux tangents and cotangents are therefore pathwise derivatives under common -random numbers. SDE/SDAE aux is saved only at actual grid nodes; -`SaveAt(ts=...)` remains unsupported because deterministic interpolation is -not valid for rough paths. - -## Deterministic interpolation and AD - -`SaveAt(t_1=True)` evaluates saved aux only at the final state. -`SaveAt(steps=True)` stores it at the initial and accepted nodes and applies -the same prefix/padding mask as the state. - -By default, `SaveAt(ts=grid)` uses normalized cubic Hermite interpolation for -ODE and root-restored DAE aux. Endpoint aux slopes are JVPs along the full -solution velocity, so they include direct parameter dependence and indirect -dependence through the state and any implicit algebraic root. Rodas5P obtains -endpoint state velocities from its published stiff-aware continuous extension, -then uses the same JVP construction for aux. No requested-time algebraic root -or aux recomputation is performed on this dense-output path. See -[Rodas5P](rodas5p.md) and [Semi-Explicit DAEs](dae.md) for the dense-output -details and links to SciML's implementation. - -For an explicit fixed-step ODE with `SaveAt(ts=grid, exact=True)`, every -requested time must instead coincide with a realized solver knot. The solver -selects those states without Hermite interpolation and evaluates aux directly -at the requested knots. Saved-aux validity is therefore checked only at those -requested knots on this exact path. - -Ordinary JAX transformations compose through every saved or interpolated aux -leaf. This includes `jax.jvp`, reverse-mode VJP/`jax.grad`, `vmap`, and -reverse-over-forward. Adaptive accept/reject decisions and mesh selection -retain the package's frozen-controller derivative convention. Consequently, -adaptive `SaveAt(steps=True)` aux derivatives hold internal times fixed and -omit mesh motion; fixed requested-grid aux does not have a moving output-time -axis. See [Adaptive Stepping and AD](adaptive_ad.md#frozen-mesh-versus-a-moving-mesh-derivative). - -## Failure behavior - -Every saved aux leaf and every inexact algebraic-aux leaf must remain finite. -Invalid algebraic context at initialization fails before time stepping. -Invalid algebraic context at a later required stage/node, or invalid saved aux -in a prefix-saving mode, terminates at the previous accepted node. -Endpoint-only saved aux is evaluated after integration; if it is invalid, the -endpoint state is retained, aux is zero-filled, and `sol.ok` is false. - -For batched differentiation where inactive lanes may leave the model domain, -`failure_ad_reference` supplies a finite point used only for safe -linearization. ODE/SDE references have the form `(x, t, p)`; DAE/SDAE -references have `(y, z, t, p)`. Without one, tinydiffeq uses an all-ones -best-effort reference. For DAE/SDAE roots, it is substituted before nlls only -after a lane is already inactive. A newly attempted root and its model context -must be JVP-safe at the actual initial point; the reference is not a post-hoc -replacement for an active failure. diff --git a/docs/dae.md b/docs/dae.md index 8aff8ec..ea5c966 100644 --- a/docs/dae.md +++ b/docs/dae.md @@ -8,36 +8,22 @@ $$ $$ where the algebraic equation is square and $g_z$ is nonsingular along the -solution. `y` and `z` may independently be array or pytree states. Leaves -within each state share one real floating dtype; the `y` and `z` dtypes may -differ. The residual `g` is a single array whose flattened size matches the -total size of `z`. The implementation supports root-restored RK4/Tsit5 and -linearly implicit Rodas5P with fixed or adaptive control. +solution. `y` and `z` may independently be array or pytree states; their +dtypes may differ. The residual `g` is a single array whose flattened size +matches the total size of `z`. Root-restored RK4/Tsit5 and linearly implicit +Rodas5P are supported, with fixed or adaptive control. The algebraic solve uses [`nlls-gram`](https://highdimensionaleconlab.github.io/nlls_gram/)'s -general Levenberg–Marquardt solver for both the primal root and its implicit -derivative. The default `Cholesky()` primal solver selects the dense normal form -for the square algebraic system, while nlls's default `ad_solver` selects the -direct nonsymmetric `LU()` rule for that square system. The implicit rule -differentiates the defining equation rather than the LM iterations; no implicit -ridge is added. Tinydiffeq applies the DAE validity policy around that solve but -does not implement a separate square-system IFT. - -Rodas5P is a JAX adaptation of Steinebach's method following SciML's -[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock) -implementation. See [Rodas5P](rodas5p.md) for direct links to SciML's tableau, -step, and interpolation sources. +Levenberg–Marquardt solver for both the primal root and its implicit +derivative: the primal defaults to dense `Cholesky()` normal equations, and +the square implicit rule defaults to a direct nonsymmetric `LU()` solve of +the defining constraint rather than a differentiation of the LM iterations. ## Minimal examples -Consider - -$$ -\dot y = pz, \qquad 0=z-y, -$$ - -whose reduced solution is $y(t)=z(t)=y_0e^{pt}$. +Consider $\dot y = pz$, $0 = z - y$, whose reduced solution is +$y(t) = z(t) = y_0 e^{pt}$. ```python import jax.numpy as jnp @@ -74,19 +60,12 @@ adaptive = solve_semi_explicit_dae( ) adaptive.aux["flow"] - -linearly_implicit = solve_semi_explicit_dae( - f, g, Rodas5P(), 0.0, 1.0, - jnp.asarray(1.0), jnp.asarray(0.5), - p=jnp.asarray(2.0), dt_0=0.1, - controller=IController(), max_steps=128, -) ``` -`z_0` is a root-finding guess, not an assumed-consistent initial value. Both +`z_0` is a root-finding guess, not an assumed-consistent initial value: both calls first solve `g(y_0, z, t_0, args, p) = 0`, so the `0.5` guess becomes -the consistent value `1.0`. RK4 and Tsit5 then solve the algebraic equation at -every stage. Rodas5P performs no further nonlinear solves. +the consistent value `1.0`. RK4 and Tsit5 then re-solve the algebraic +equation at every stage. ## Nonlinear-solve and AD contract @@ -97,7 +76,6 @@ from tinydiffeq import LMRootSolver root_solver = LMRootSolver( max_steps=8, - max_steps_is_success=False, # compatibility field; MAX_STEPS stays invalid atol=None, # 1e-6 float32, 1e-10 float64 gtol=0.0, # required: residual stopping only xtol=0.0, # required: residual stopping only @@ -106,30 +84,16 @@ root_solver = LMRootSolver( ) ``` -The outer `max_steps` counts attempted time steps, including adaptive -rejections. `root_solver.max_steps` separately bounds one algebraic root. For -Rodas5P it affects only initial consistency; the method's later stages reuse -one dense LU factorization per attempted time step. - -Root tolerances are independent of the outer controller tolerances. Explicit -`atol` must be positive; `None` selects the dtype default, and `atol=0` is -invalid. `gtol` and `xtol` must both equal zero; `LMRootSolver` rejects nonzero -values so `CONVERGED` can only come from the residual test. Every accepted -algebraic root must report `CONVERGED` and have Euclidean residual norm -`sqrt(sum(residual**2))` strictly below the root `atol`. -`max_steps_is_success` remains in the configuration for source compatibility, -but Tinydiffeq always asks nlls to treat `MAX_STEPS` as a failed implicit solve; -setting the field to `True` does not broaden DAE root acceptance. - -The primal nonlinear solve uses nlls-gram's dense `Cholesky()` normal-equation -default. For a successful square root, nlls's implicit rule forms `dg/dz`, -forms the right-hand side with respect to `(y, t, p)`, and applies its direct -`LU()` square solve. Transposing that nlls rule supplies the VJP. The primal LM -factorization is not reused because it represents a damped normal equation, -not the accepted root's generally nonsymmetric Jacobian. A failed nlls status -has zero implicit tangent by the nlls solve contract. For the rare root solve -that needs to depart from the defaults, `solver_options` forwards constructor -arguments for either the primal `linear_solver` or the implicit `ad_solver`: +The outer `max_steps` counts attempted time steps; `root_solver.max_steps` +separately bounds one algebraic root, and root tolerances are independent of +the controller tolerances. `gtol` and `xtol` must be zero so `CONVERGED` can +only come from the residual test: every accepted root must report +`CONVERGED` with Euclidean residual norm strictly below `atol`, and a +`MAX_STEPS` iterate is never a differentiable root. + +`solver_options` forwards constructor arguments verbatim to nlls-gram's +`LevenbergMarquardt` — either the primal `linear_solver` or the implicit +`ad_solver`: ```python from nlls_gram import QR @@ -138,100 +102,111 @@ 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 -those two options rather than silently honoring an override. `ad_solver` -remains nlls-owned and may be supplied explicitly; the square default is -`LU()`. - -`predictor="previous"` is the default: each explicit RK stage starts from the -most recent successful algebraic root. `predictor="secant"` extrapolates from -the accepted-step root through the most recent successful stage at a later -time. Duplicate RK4 stage times, non-forward targets, and failed stages fall -back to the previous root. The predictor is stop-gradiented, so a successful -root still uses the same implicit derivative. Its time-derived extrapolation -scale is cast separately to each algebraic leaf's dtype, preserving the `z` -dtype when the differential/time and algebraic dtypes differ. Secant prediction -assumes the continued branch is locally unique; with multiple roots or a tight -finite iteration budget, changing the guess can change the selected branch, -value, or status. - -Every nonlinear root passes `(y, t, p)` to nlls-gram. Its implicit rule then +than being mirrored here. Options are normalized to a sorted tuple so equal +configurations stay hashable and share one compiled solver. +`cache_jacobian` and `geodesic_acceleration` are fixed to `False` and +rejected — each DAE stage changes the root problem, and the intended path is +the ordinary dense LM step. + +`predictor="previous"` (default) starts each explicit RK stage from the most +recent successful root. `predictor="secant"` extrapolates from the +accepted-step root through the most recent successful stage at a strictly +later time; duplicate stage times and failed stages fall back to the +previous root. Predictor values are stop-gradiented, so successful roots keep +the same implicit derivative — but with multiple algebraic roots a different +warm start can select a different branch, so secant assumes the continued +root is locally unique. + +Every nonlinear root passes `(y, t, p)` to nlls-gram, whose implicit rule differentiates the defining constraint, $$ -\dot z = -g_z^{-1} - (g_y\dot y + g_t\dot t + g_p\dot p), +\dot z = -g_z^{-1}(g_y\dot y + g_t\dot t + g_p\dot p), $$ -rather than differentiating the LM iterations. The warm-start guess has zero -derivative by design. Rodas5P differentiates through its exact JAX Jacobian, -time derivative, LU factorization, and linear stage solves. `args` is fixed -data; put every differentiated model quantity in `p`. On the default bounded -integration path, JVP, VJP, `vmap`, and reverse-over-forward compose through -the complete DAE solve. - -`sol.num_steps` counts logical time-step attempts, including adaptive -rejections. `sol.num_root_solves` counts active nonlinear root calls, including -the initial consistency solve and failed calls, and `sol.num_root_steps` sums -their LM update counts. Rodas5P therefore reports one root solve regardless of -its time-step count; later linear stages are not nonlinear roots. These -counters have exact-zero tangents. Under `vmap`, a masked lane may still execute -physically while remaining absent from its logical counters. - -The default `adaptive_loop="bounded"` uses a reverse-mode-capable bounded scan. -`adaptive_loop="forward"` uses a dynamic actual-work loop for adaptive Tsit5 -and Rodas5P. It supports primal evaluation, JVP, and nested forward mode, but -not ordinary reverse mode. A vmapped forward loop runs until the slowest lane -finishes. - -The differential field may return `(dy, saved_aux)`. Saved aux is a nonempty -pytree of nonempty real floating arrays; different leaves may use different -floating dtypes. tinydiffeq evaluates it at required saved nodes. Ordinary JAX -differentiation composes with either the root's implicit derivative or the -Rodas5P stages, so aux tangents and cotangents include both direct dependence -on `p` and indirect dependence through `z`. - -The algebraic function may instead or additionally return -`(residual, algebraic_aux)`. In that case `f` takes -`(y, z, t, args, p, algebraic_aux)`. This value is internal cached context: -the nonlinear solver sees only the residual, and only differential-field -saved aux appears in `sol.aux`. See [Auxiliary Outputs](aux.md) for the four -supported combinations and flag behavior. - -Every saved aux leaf and every inexact algebraic-aux leaf must be finite. -Invalid algebraic context at initialization sets `ok=False` before any -time-step work. `SaveAt(steps=True)` and `SaveAt(ts=...)` check saved aux at -the initial and accepted nodes, so an invalid value freezes the previous valid -prefix. Endpoint mode evaluates saved aux only after integration; an invalid -final value retains the endpoint state, returns zero aux, and sets `ok=False`. - -An adaptive stage-root failure rejects the time-step attempt and asks the -controller for a smaller step; a fixed-step failure terminates. Rodas5P linear -failures follow the same controller policy. -In either case `sol.ok` is false if the endpoint is not reached with valid -algebraic states. nlls supplies the primal LM iterate, diagnostics, and the -implicit JVP/VJP. Tinydiffeq applies the DAE's residual-and-status acceptance -check and returns the differentiation-inert warm-start guess when that check -fails; aux at a failed initial root is a zero pytree of the declared shape. -Callers that want to retain successful-lane JVPs/VJPs after another lane has -already become inactive should pass -`failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)`, choosing a point where -the residual, context, and saved-aux maps are finite and differentiable. -tinydiffeq substitutes this point into an already-inactive root call before -entering nlls, and also uses it for inactive algebraic aux, saved aux, and -differential-field evaluations. It is not an nlls solve argument and never -changes an active root attempt. - -Every newly attempted root must therefore be JVP-safe at its actual -`(y, z_guess, t, p)`: the residual and any model context evaluated there must -have valid derivatives. The reference cannot rescue an intrinsically invalid -active attempt after the fact. Without an explicit reference, an all-ones -best-effort default is used for inactive work; gradients are not guaranteed if -the model is undefined there. A failed lane itself is never a valid solution. +rather than the LM iterations; the warm-start guess has zero derivative by +design. `args` is fixed data; put every differentiated quantity in `p`. On +the default bounded path, JVP, VJP, `vmap`, and reverse-over-forward compose +through the complete solve. `adaptive_loop="forward"` runs an actual-work +loop for adaptive Tsit5 and Rodas5P (primal, JVP, and nested forward mode +only), as for [ODEs](ode.md#static-shapes-and-saveat). + +`sol.num_steps` counts logical time-step attempts including rejections; +`sol.num_root_solves` counts active nonlinear root calls (including the +initial consistency solve and failures) and `sol.num_root_steps` sums their +LM updates. All have exact-zero tangents. An adaptive stage-root failure +rejects the attempt and retries with a smaller step; a fixed-step failure +terminates. Either way `sol.ok` is false if `t_1` is not reached with valid +algebraic states. + +For batched differentiation where a lane may fail or leave the model domain, +pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)` at a point where +the residual, context, and aux maps are finite and differentiable. It is +substituted into already-inactive root calls and inactive aux/field +evaluations so masked lanes cannot poison the JVP/VJP of successful lanes. A +newly attempted root must still be JVP-safe at its actual +`(y, z_guess, t, p)`; the reference never rescues an active failure, and a +failed lane's outputs are not a valid solution. Without an explicit +reference, an all-ones best-effort default is used. + +## Auxiliary outputs + +Saved output belongs to the differential field, and the algebraic function +may separately expose internal context: + +```python +def g(y, z, t, args, p): + return residual, algebraic_aux + + +def f(y, z, t, args, p, algebraic_aux): + return dy, saved_aux +``` + +All four combinations are supported: neither aux, saved aux only, algebraic +aux only, or both. When algebraic aux is present, `f` must take the full +six-argument form. `has_aux` / `has_algebraic_aux` default to abstract +auto-detection; explicit `False` skips those traces. + +Algebraic aux is internal cached context: it is passed to `f`, included in +the implicit derivative path, ignored by the nonlinear solver's residual +interface, and never stored or interpolated. It may be a nonempty pytree of +bool, integer, real, or complex arrays; every inexact leaf must be finite. +Only the differential field's real-floating `saved_aux` becomes `sol.aux`, +stored at accepted nodes and interpolated on requested grids. Invalid +algebraic context at initialization fails before any time-step work; invalid +saved aux in a prefix mode terminates at the previous accepted node, while +endpoint mode keeps the endpoint state with zero aux and `ok=False`. + +## Rodas5P for DAEs + +For the stiff path, tinydiffeq constructs the flattened mass-matrix system +internally: + +$$ +M\dot u=F(u,t), \qquad +u=(y,z), \quad M=\operatorname{diag}(I_y,0_z), \quad F=(f,g). +$$ + +```python +linearly_implicit = solve_semi_explicit_dae( + f, g, Rodas5P(), 0.0, 1.0, + jnp.asarray(1.0), jnp.asarray(0.5), + p=jnp.asarray(2.0), dt_0=0.1, + controller=IController(), max_steps=128, +) +``` + +`LMRootSolver` is used once, for initial consistency. Every later Rodas5P +stage solves $\left(M/(\gamma h) - F_u\right)k_i = r_i$ with one reused +pivoted LU factorization per attempted step — no nonlinear endpoint +restoration. Returned internal `z` values therefore satisfy the constraint +to the method's integration accuracy, not to `LMRootSolver.atol`, and the +solver reports one root call regardless of its attempt count. This +intentionally differs from RK4/Tsit5, which root-solve every stage; Rodas5P +is the choice when those algebraic solves dominate runtime or the coupled +dynamics are stiff. See [Stiff ODEs: Rodas5P](ode.md#stiff-odes-rodas5p) for +the method, its AD boundaries, and the SciML/Steinebach credit and links. ## Saving output @@ -240,79 +215,47 @@ All `SaveAt` modes are supported: - `SaveAt(t_1=True)` returns the endpoint. - `SaveAt(steps=True)` returns the initial point and accepted internal steps as a padded `max_steps + 1` buffer with the usual `accepted` mask. -- `SaveAt(ts=grid)` uses cubic Hermite for root-restored methods and Rodas5P's - stiff-aware continuous extension for `(y, z)`. Aux uses cubic Hermite in - both cases. It performs no query-time nonlinear solves. - -`SaveAt(ts=..., exact=True)` is not a DAE mode; exact knot gathering is limited -to explicit fixed-step ODEs. +- `SaveAt(ts=grid)` uses cubic Hermite for root-restored methods and + Rodas5P's stiff-aware continuous extension for `(y, z)`; aux uses cubic + Hermite in both cases. No query-time nonlinear solves are performed. + `exact=True` is not a DAE mode. The result is a `DAESolution` with `ts`, `ys`, `zs`, `ok`, `num_accepted`, -`accepted`, `aux`, `num_steps`, `num_root_solves`, and `num_root_steps` fields. -For pytree states, saved rows are a leading axis on every state and aux leaf; -the one `accepted` mask applies to the complete output. +`accepted`, `aux`, `num_steps`, `num_root_solves`, and `num_root_steps`. ### Dense output for root-restored RK4 and Tsit5 At a consistent knot, differentiating the constraint gives - -$$ -g_z\dot z = -(g_y\dot y + g_t). -$$ - -tinydiffeq solves this linear system once per accepted knot only when a query -grid is requested. It then obtains `aux_dot` by a JVP of the aux map along -$(\dot y,\dot z,1)$. Values and total derivatives feed the same normalized -cubic Hermite basis used for ODE states. This is an order-3 continuous -extension—uniform interpolation error $O(h^4)$—when `f` and `g` are $C^4$, -$g_z$ stays uniformly nonsingular near the solution, and root error is no -larger than the desired dense-output error. RK4 and Tsit5 knot errors meet the -required order under their usual assumptions. - -The normalized coordinate stays in `[0, 1]` and Hermite basis coefficients -are bounded by 3, which is favorable in float32. SciML's specialized Tsit5 -dense polynomial has one higher order for `y`, but requires all seven stages, -does not directly supply `z`/aux output, and has much larger coefficients. -Using one Hermite construction keeps `y`, `z`, and aux at the same dense order -with substantially less storage. - -Interpolated `z` and aux are approximations: away from accepted knots they -need not satisfy `g=0` exactly. The constraint defect is $O(h^4)$ under the -conditions above. Use `SaveAt(steps=True)` when every returned row must be an -actual converged root. Dense output also requires one `g_z` factorization per -accepted knot, rather than one nonlinear solve per requested time; its cost -therefore scales with internal steps rather than grid length. +$g_z\dot z = -(g_y\dot y + g_t)$. tinydiffeq solves this linear system once +per accepted knot (only when a query grid is requested), obtains `aux_dot` +by a JVP of the aux map along $(\dot y, \dot z, 1)$, and feeds values and +derivatives into the same normalized cubic Hermite basis used for ODE +states — an order-3 continuous extension with uniform error $O(h^4)$ when +`f` and `g` are $C^4$, $g_z$ stays uniformly nonsingular, and root error is +below the dense-output error. Interpolated `z` and aux are approximations: +away from knots the constraint defect is $O(h^4)$. Use `SaveAt(steps=True)` +when every returned row must be an actual converged root. ### Dense output for Rodas5P -Rodas5P stores the three coefficient pytrees defined by Steinebach's -fourth-order stiff-aware continuous extension. tinydiffeq evaluates the same -polynomial form used by +Rodas5P stores the three coefficient pytrees of Steinebach's fourth-order +stiff-aware continuous extension and evaluates the same polynomial form as [SciML's Rosenbrock interpolant](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_interpolants.jl) -for the combined `(y, z)` state. No `g_z` factorization or nonlinear solve is -performed for requested times. - -Aux remains a stored accepted-knot quantity. Its cubic-Hermite endpoint -tangents come from the Rodas polynomial's endpoint derivatives and a JVP of -the aux map. Aux is therefore interpolated rather than recalculated at every -query. Rodas5P accepted knots are not root-restored: their constraint defect, -and that of dense output, is controlled by integration accuracy rather than -`LMRootSolver.atol`. - -Knot selection and adaptive step sizes remain differentiation-inert under the -frozen-controller convention. Values, implicit slopes, and aux differentiate -on that realized mesh, but adaptive `SaveAt(steps=True)` omits mesh motion. -See [Frozen mesh versus a moving-mesh derivative](adaptive_ad.md#frozen-mesh-versus-a-moving-mesh-derivative). -If `sol.ok` is false, neither outputs nor their derivatives should be treated -as a valid solution. +for the combined `(y, z)` state — no $g_z$ factorization or nonlinear solve +at query times. Aux endpoint tangents come from the Rodas polynomial's +endpoint derivatives and a JVP of the aux map. + +Knot selection and step sizes are differentiation-inert under the +frozen-controller convention; see +[AD through adaptive stepping](ode.md#ad-through-adaptive-stepping). If +`sol.ok` is false, neither outputs nor derivatives are a valid solution. ## Deliberate limits Only the internally constructed constant block mass matrix -`diag(I_y, 0_z)` is supported; there is no public general mass-matrix or fully -implicit residual API. Rodas5P uses dense Jacobians and dense pivoted LU, not -sparse or Krylov linear algebra. Higher-index constraints and automatic index -reduction are unsupported. This is an initial-value solver: it does not -determine unknown initial costates or solve boundary-value or saddle-path -conditions. Initial branch selection and jumps between multiple roots are not -differentiable. +`diag(I_y, 0_z)` is supported; there is no general mass-matrix or fully +implicit residual API. Rodas5P uses dense Jacobians and dense pivoted LU. +Higher-index constraints and automatic index reduction are unsupported. This +is an initial-value solver: it does not determine unknown initial costates +or solve boundary-value problems, and jumps between multiple root branches +are not differentiable. diff --git a/docs/index.md b/docs/index.md index 2b74c08..9a19dfc 100644 --- a/docs/index.md +++ b/docs/index.md @@ -1,64 +1,25 @@ # tinydiffeq -`tinydiffeq` is a deliberately tiny set of ODE/SDE/DAE/SDAE integrators and -finite-state Markov simulators for JAX: fixed-step Euler and RK4, adaptive Tsit5 with integral or -proportional-integral step-size control, linearly implicit Rodas5P for stiff -ODEs and index-1 DAEs, and fixed-step Euler–Maruyama for Itô SDEs and SDAEs. -Fixed stepping and the default adaptive path use bounded `lax.scan` loops with -static shapes. These solves support forward mode, reverse mode, and -reverse-over-forward. Adaptive ODE and DAE solves also offer -`adaptive_loop="forward"`, a dynamic actual-work loop for primal, JVP, and -nested forward AD; JAX cannot reverse-transpose this path. -Finite-state DTMC/CTMC simulation is primal-only and offers chronological scan -and associative parallel-prefix methods. Fixed-chain probability forecasts use -binary matrix powers for DTMC endpoints, dense exponentials for small CTMCs, or -fixed or adaptive matrix-free Arnoldi/Krylov actions over array or pytree -probabilities; see -[Markov Chains](markov_chains.md). The same backends solve general fixed -homogeneous linear systems; see [Linear Exponential Solves](exponential.md). - -Rodas5P follows SciML's -[`OrdinaryDiffEqRosenbrock`](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock) -implementation and Steinebach's published method. **Use SciML or +`tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE/DAE/SDAE +integrators and finite-state Markov simulators for JAX: fixed-step Euler and +RK4, adaptive Tsit5, linearly implicit Rodas5P for stiff ODEs and index-1 +DAEs, and fixed-step Euler–Maruyama, Milstein, and SRA1 for Itô SDEs and +SDAEs. Solves run in bounded `lax.scan` loops with static shapes and support +forward mode, reverse mode, and reverse-over-forward; adaptive ODE/DAE solves +can opt into a dynamic actual-work loop (`adaptive_loop="forward"`, no +reverse mode). Probability forecasts and general fixed homogeneous linear +solves use matrix powers, dense exponentials, or matrix-free Krylov actions; +see [Markov Chains](markov_chains.md) and +[Linear Exponential Solves](exponential.md). + +**Use [SciML](https://docs.sciml.ai/DiffEqDocs/stable/) or [diffrax](https://docs.kidger.site/diffrax/) instead if you need any of:** - general mass matrices, fully implicit solvers, or higher-index DAEs -- sparse/Krylov linear solves and preconditioners inside general ODE/DAE stages -- full derivative-term PID step-size control +- sparse/Krylov linear solves and preconditioners inside ODE/DAE stages +- adaptive SDE stepping (Brownian-bridge noise), full PID step-size control - events, root-finding, or backward-time integration -- dense output / continuous interpolation objects -- checkpointed or backsolve adjoints for long horizons - -Use the default bounded loop when reverse mode is required. Use the forward -loop when the actual adaptive attempt count is much smaller than `max_steps` -and the caller needs only primal or forward-mode execution. Both retain static -public output shapes; under `vmap`, a dynamic loop still runs until the slowest -lane finishes. - -## 2.4.0 migration note - -- `SaveAt(ts=..., exact=True)` now gathers realized knots for explicit - fixed-step ODEs. Every query must align with a knot; adaptive methods, - Rodas5P, DAEs, SDEs, and SDAEs continue to reject exact mode. -- `Solution.num_steps` and `DAESolution.num_steps` count logical attempts, - including rejections. DAE results additionally expose `num_root_solves` and - `num_root_steps`; `num_accepted` retains its existing meaning. -- Adaptive ODE and DAE solves may opt into `adaptive_loop="forward"` for an - actual-work loop. It supports primal, JVP, and nested forward AD but not - reverse mode; `adaptive_loop="bounded"` remains the reverse-mode-capable - default. Under `vmap`, the forward loop runs to the slowest lane. -- `LMRootSolver(predictor="secant")` is an opt-in continuation warm start for - locally unique algebraic branches; `predictor="previous"` remains the - default. - -DAE root acceptance is stricter in 2.4.0. nlls-gram owns both the primal root -solve and implicit derivative; square implicit AD defaults to direct `LU()`. -Only `CONVERGED` roots whose residual norm is below `atol` are accepted, so -`gtol` and `xtol` must both be zero. `max_steps_is_success` remains for source -compatibility, now defaults to `False`, and never makes `MAX_STEPS` a valid -root. Upgrading configurations should remove nonzero `gtol`/`xtol`; if they -relied on budget exhaustion, increase the root budget or adjust the residual -tolerance instead. +- dense output objects or checkpointed/backsolve adjoints for long horizons ## Install @@ -86,26 +47,17 @@ f(x, t, args, p) ``` - `x` is an array or pytree state. Leaves must share one real floating dtype; - the vector field returns the same structure and dtype. -- `t` is time. -- `args` is pass-through data. By convention it is **not** an AD target — - nothing stops you differentiating with respect to it, but the library's - contracts and tests treat it as constants. -- `p` holds differentiable parameters — any pytree, e.g. neural-network - weights. jvp/vjp with respect to `p` and `x_0` are first-class and tested. - -The arity is inspected once and the function is wrapped into the canonical -four-argument form, so the compiled code is identical for all four. There is -no special autonomous code path: an unused `t` is dead-code-eliminated. -`drift` and `diffusion` in [`solve_sde`](sde.md) follow the same convention. -Semi-explicit DAE fields use `(y, z)`, `(y, z, t)`, -`(y, z, t, args)`, or `(y, z, t, args, p)`; see -[Semi-Explicit DAEs](dae.md). -Differential fields and stochastic drifts may return `(value, saved_aux)`; -DAE/SDAE algebraic functions may separately return internal context consumed -by the dynamics. Saved aux is differentiated and follows `SaveAt`. See -[Auxiliary Outputs](aux.md) for all supported contracts and -[Semi-Explicit SDAEs](sdae.md) for the stochastic form. + the field returns the same structure and dtype. +- `args` is pass-through data — by convention **not** an AD target. +- `p` holds differentiable parameters (any pytree, e.g. network weights). + JVP/VJP with respect to `p` and `x_0` are first-class and tested. + +The arity is inspected once and wrapped into the canonical four-argument +form, so the compiled code is identical for all four. `drift` and `diffusion` +in [`solve_sde`](sde.md) follow the same convention; semi-explicit DAE fields +use `(y, z)` through `(y, z, t, args, p)` — see +[Semi-Explicit DAEs](dae.md) and [SDAEs](sdae.md). Fields may return +`(value, saved_aux)` to save extra quantities with the solution. ## Minimal example @@ -133,12 +85,6 @@ sol.xs # (21,) states on the grid, however many internal steps were taken sol.ok # False if integration or a requested output failed ``` -With no arguments, `IController()` and `PIController()` use precision-aware -tolerances: `rtol=1e-4, atol=1e-6` for float32 states and -`rtol=1e-7, atol=1e-9` for float64 states. Explicit values override the -policy and are cast to the state dtype. Automatic `dt_min` is -`10 * eps * max(1, abs(t_1))` in the time dtype. - Gradients go straight through the solve: ```python @@ -151,45 +97,31 @@ def endpoint(p): jax.grad(endpoint)(jnp.asarray(1.3)) # reverse mode jax.jvp(endpoint, (jnp.asarray(1.3),), (jnp.asarray(1.0),)) # forward mode -jax.grad(lambda p: jax.jvp(endpoint, (p,), (jnp.asarray(1.0),))[1])( - jnp.asarray(1.3) -) # reverse-over-forward ``` ## Design contracts at a glance - **`dt_0` is required.** There is no initial-step heuristic. -- **`max_steps` counts attempted internal steps**, including rejections. It - controls the bounded scan and only becomes an output-row count in - `SaveAt(steps=True)`, which returns `max_steps + 1` padded rows including - the initial state. Accepted steps form a contiguous prefix; rejected - attempts are omitted and the tail repeats the last accepted state. - `sol.num_steps` reports actual attempts and `sol.num_accepted` reports - successful advances. +- **`max_steps` counts attempted internal steps**, including rejections. + `sol.num_steps` reports attempts, `sol.num_accepted` successful advances. +- **`SaveAt` is the shape contract**: endpoint, fixed interpolation grid, or + padded accepted-step prefix — output shapes never depend on how many steps + the controller took. See [ODEs](ode.md#static-shapes-and-saveat). - **Fixed-step times do not depend on the attempt budget.** They are formed - arithmetically from the accepted-step index, with only a small local endpoint - snap. Increasing a nonbinding `max_steps` therefore does not change the - numerical method. -- **Exact fixed-step output is explicit-ODE-only.** - `SaveAt(ts=grid, exact=True)` gathers internal knots without interpolation; - every query must align with a realized knot. -- **Adaptive loop choice is an AD choice.** `adaptive_loop="bounded"` is the - reverse-mode-capable default. `adaptive_loop="forward"` executes actual - attempts but supports only primal, JVP, and nested forward mode. + arithmetically from the accepted-step index. +- **The controller is stop-gradiented.** States differentiate through solver + stages on the realized, frozen mesh; see + [AD through adaptive stepping](ode.md#ad-through-adaptive-stepping). - **Forward time only**: `t_1 > t_0`. -- **Never poisons.** `sol.ok` reports whether `t_1` was reached and every - requested output was valid; callers that want diverging residuals can map - `jnp.where(sol.ok, x, jnp.inf)` over `sol.xs`. -- **`project`** (an idempotent clamp, e.g. positivity) is applied at every - point where the vector field is evaluated and to every accepted state. -- **Never sets `jax_enable_x64`.** The time dtype follows the common state - dtype; float32 problems stay float32 even when x64 is enabled. +- **Never poisons.** `sol.ok` reports failure; callers map + `jnp.where(sol.ok, x, jnp.inf)` when they want loud divergence. +- **Never sets `jax_enable_x64`.** The time dtype follows the state dtype; + float32 problems stay float32 even when x64 is enabled. - Solvers, controllers, `SaveAt`, and `Solution` are frozen dataclasses - registered as pytrees: numeric fields (tolerances, grids, `dt_0`, `x_0`) are - data leaves, so changing them never recompiles. - -Read next: [Static Shapes](static_shapes.md) for the loop and -`SaveAt`, [Adaptive Stepping and AD](adaptive_ad.md) for what is and is not -differentiated, [Auxiliary Outputs](aux.md), [Rodas5P](rodas5p.md) for the SciML-derived linearly implicit -method, [DAEs](dae.md), [SDEs](sde.md), [SDAEs](sdae.md), and the [API -Reference](api.md). + registered as pytrees: numeric fields (tolerances, grids, `dt_0`, `x_0`) + are data leaves, so changing them never recompiles. + +Read next: [ODEs](ode.md), [SDEs](sde.md), [Semi-Explicit DAEs](dae.md), +[SDAEs](sdae.md), [Markov Chains](markov_chains.md), +[Linear Exponential Solves](exponential.md), and the +[API Reference](api.md). diff --git a/docs/llms.txt b/docs/llms.txt index b2cb761..72362da 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -1,17 +1,14 @@ # tinydiffeq -> Tiny ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Fixed stepping and the default adaptive path use bounded lax.scan loops with static output shapes and forward/reverse AD, including reverse-over-forward. Adaptive Tsit5 and Rodas5P ODE/DAE solves can instead select adaptive_loop="forward", an actual-work lax.while_loop for primal, JVP, and nested forward AD; reverse mode is unsupported, and vmap runs to the slowest lane. max_steps is an attempt budget, while Solution.num_steps records actual attempts. Fixed-step times are arithmetic and budget-invariant. SaveAt(ts=..., exact=True) gathers aligned internal knots for explicit constant-step ODEs only; other requested-grid deterministic output uses Hermite or Rodas5P interpolation. Adaptive SaveAt(steps=True) AD freezes the internal mesh and omits mesh motion, an asymptotically exact residual Jacobian only near a smooth residual root. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, an embedded estimator, and its published stiff-aware fourth-order dense output. RK4/Tsit5 DAEs delegate both their stage roots and implicit derivatives to nlls-gram; the square primal defaults to normal Cholesky and the square implicit rule defaults to direct nonsymmetric LU. LMRootSolver requires residual-only stopping with gtol=xtol=0: every accepted algebraic root must report CONVERGED and have Euclidean residual norm below root atol. max_steps_is_success remains for source compatibility but never makes MAX_STEPS a valid DAE root. LMRootSolver defaults to the previous-root predictor and offers an opt-in secant predictor for locally unique branches. DAESolution reports logical time attempts, nonlinear root calls, and LM update counts; Rodas5P uses LM only for initial consistency. ODE fields and stochastic drifts may return (value, saved_aux), while DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary JAX pytrees. IController remains the adaptive default; PIController is opt-in. Finite-state DTMC/CTMC sampling is primal-only and offers sequential and associative parallel-prefix methods. Deterministic distribution forecasts are differentiable in their initial mass through matrix powers, dense exponentials, or matrix-free Arnoldi/Krylov actions. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints. +> Tiny differentiable ODE/SDE/DAE/SDAE solvers and finite-state Markov tools for JAX. Fixed stepping and the default adaptive path use bounded lax.scan loops with static output shapes and forward/reverse AD, including reverse-over-forward; adaptive Tsit5 and Rodas5P ODE/DAE solves can instead select adaptive_loop="forward", an actual-work lax.while_loop for primal, JVP, and nested forward AD only. max_steps is an attempt budget, Solution.num_steps records actual attempts, and fixed-step times are arithmetic and budget-invariant. SaveAt picks the endpoint, a fixed interpolation grid (exact=True gathers aligned knots for explicit constant-step ODEs), or the padded accepted-step prefix. Adaptive AD freezes the internal mesh (stop-gradiented controllers) and omits mesh motion. Rodas5P supports stiff ODEs and semi-explicit index-1 DAEs with exact dense JAX Jacobians, one reused LU factorization per attempt, and Steinebach's stiff-aware fourth-order dense output, following SciML's OrdinaryDiffEqRosenbrock. RK4/Tsit5 DAEs delegate stage roots and implicit derivatives to nlls-gram; LMRootSolver requires residual-only stopping (gtol=xtol=0) and accepts only CONVERGED roots with Euclidean residual norm below the root atol, with previous-root or opt-in secant stage predictors. solve_sde integrates diagonal-noise Ito SDEs with EulerMaruyama (strong order 0.5), Milstein (1.0, commutative diagonal), or SRA1 (1.5, additive); each solver declares its per-step noise via sample_noise, drawn from a key or passed explicitly through noise= as validated, differentiable data. solve_semi_explicit_sdae applies EulerMaruyama or SRA1 to the reduced index-1 stochastic system with root-restored consistency; SaveAt(ts=...) raises for SDEs/SDAEs because interpolating rough paths is wrong. Fields and stochastic drifts may return (value, saved_aux); DAE/SDAE algebraic functions may return internal (residual, algebraic_aux) context. States may be arrays or arbitrary pytrees of same-dtype real floating arrays, and the library never sets jax_enable_x64. Finite-state DTMC/CTMC sampling is primal-only; deterministic distribution forecasts are differentiable in their initial mass through matrix powers, dense exponentials, or matrix-free Arnoldi/Krylov actions. Use SciML/diffrax for general mass matrices, fully implicit or higher-index DAEs, adaptive stochastic stepping, events, continuous solution objects, or specialized adjoints. ## Docs -- [Home — positioning, vector-field signature convention f(x, t, args, p), minimal examples](https://highdimensionaleconlab.github.io/tinydiffeq/) -- [Static shapes — bounded and actual-work loops, budget-invariant fixed times, SaveAt modes, work counters](https://highdimensionaleconlab.github.io/tinydiffeq/static_shapes/) -- [Adaptive stepping and AD — frozen-mesh derivative contract, loop-mode AD boundary, controller and horizon logic](https://highdimensionaleconlab.github.io/tinydiffeq/adaptive_ad/) -- [Auxiliary outputs — saved versus algebraic aux, interpolation, AD, and failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/aux/) -- [Rodas5P — Steinebach method, direct SciML implementation links, stiff ODE/DAE formulation, dense output, and AD](https://highdimensionaleconlab.github.io/tinydiffeq/rodas5p/) -- [Semi-explicit DAEs — index-1 contract, algebraic LM roots and predictors, counters, implicit AD, failure behavior, and SaveAt](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) -- [SDEs — Euler-Maruyama orders, fixed-noise key semantics, shared-path strong-convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/) -- [Semi-explicit SDAEs — reduced-SDE Euler-Maruyama, algebraic roots, aux, convergence assumptions, and pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) +- [Home — positioning, install, vector-field signature convention f(x, t, args, p), minimal example, design contracts](https://highdimensionaleconlab.github.io/tinydiffeq/) +- [ODEs — solver table, fixed vs adaptive control, static shapes and SaveAt, AD through adaptive stepping, aux outputs, stiff Rodas5P with SciML credit](https://highdimensionaleconlab.github.io/tinydiffeq/ode/) +- [SDEs — EulerMaruyama/Milstein/SRA1 orders and noise contracts, explicit noise= as differentiable data, fixed-key semantics, shared-path convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/) +- [Semi-explicit DAEs — index-1 contract, nlls-gram LM roots and predictors, implicit AD, aux contracts, Rodas5P mass-matrix path, dense output, failure behavior](https://highdimensionaleconlab.github.io/tinydiffeq/dae/) +- [Semi-explicit SDAEs — reduced-SDE EulerMaruyama and SRA1, algebraic roots, aux, convergence assumptions, pathwise AD](https://highdimensionaleconlab.github.io/tinydiffeq/sdae/) - [Finite-state Markov chains — sampling, deterministic PMF forecasts, matrix-free Krylov CTMC actions, pytrees, vmap, and AD scope](https://highdimensionaleconlab.github.io/tinydiffeq/markov_chains/) - [Linear exponential solves — dense expm and matrix-free Arnoldi/Krylov actions for fixed linear array or pytree operators, with traced and hand-coded initial-state JVP/VJP](https://highdimensionaleconlab.github.io/tinydiffeq/exponential/) -- [API reference — ODE/SDE/DAE solves, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/) +- [API reference — solve functions, solvers, controllers, root configuration, SaveAt, solution types, interpolation, and quadrature](https://highdimensionaleconlab.github.io/tinydiffeq/api/) diff --git a/docs/ode.md b/docs/ode.md new file mode 100644 index 0000000..3887021 --- /dev/null +++ b/docs/ode.md @@ -0,0 +1,223 @@ +# ODEs + +`solve_ode` integrates `dx/dt = f(x, t, args, p)` from `t_0` to `t_1 > t_0`. +The state is an array or pytree whose leaves share one real floating dtype; +`args` is pass-through data (not an AD target by convention) and `p` holds +differentiable parameters. `dt_0` is required — there is no initial-step +heuristic. + +## Solvers + +| Solver | Order | Stepping | +|---|---|---| +| `Euler()` | 1 | fixed only | +| `RK4()` | 4 | fixed only | +| `Tsit5()` | 5(4) | fixed or adaptive (embedded estimate, FSAL) | +| `Rodas5P()` | 5(4) | fixed or adaptive, linearly implicit (stiff) | + +Fixed stepping uses `ConstantStepSize()` (the default controller): +`dt_0 = (t_1 - t_0)/n` with `max_steps = n` reproduces a uniform grid +exactly. Times are formed arithmetically as `t_0 + i * dt_0` with a small +endpoint snap that never scales with `max_steps`, so increasing a nonbinding +attempt budget cannot change the numerical method. `unroll=` (a static int, +fixed stepping only) unrolls that many steps per scan iteration — identical +values, fewer GPU dispatches, more compile time; on small-batch +neural-network fields it cut reverse-mode solve time 2–3× on an L40S (see +`benchmarks/results/`). + +Adaptive stepping uses `IController()` or `PIController()` with an embedded +error estimate. Omitted tolerances follow the state precision: + +| State dtype | `rtol` | `atol` | unit-scale `dt_min` | +|---|---:|---:|---:| +| float32 | `1e-4` | `1e-6` | `10 * eps` ≈ `1.19e-6` | +| float64 | `1e-7` | `1e-9` | `10 * eps` ≈ `2.22e-15` | + +Explicit values override the policy and are cast to the state dtype; use them +whenever tolerances are part of a reproducibility contract. Every attempt is +clipped to the remaining horizon, and the controller's next-step proposal is +computed from the clipped step — the clip doubles as the growth guard that +keeps near-flat fields from growing steps into quarter-horizon leaps. +`PIController` additionally damps step-size oscillations through the previous +accepted error ratio; `p_coeff=0, i_coeff=1` reproduces `IController` exactly. + +## Static shapes and `SaveAt` + +`max_steps` bounds attempted internal steps, including rejections. The +default `adaptive_loop="bounded"` runs a bounded `lax.scan` with exactly +`max_steps` attempt slots grouped into chunks that are skipped once the solve +completes; shapes are static, and changing tolerances, curvature, or initial +conditions never recompiles (pinned by `tests/test_recompile.py`). +`adaptive_loop="forward"` instead runs a dynamic `lax.while_loop` that +executes only actual attempts; it supports primal evaluation, JVP, and nested +forward mode, but JAX cannot transpose it, so reverse mode requires the +bounded loop. Under `vmap`, both run until the slowest lane finishes: the +bounded loop's skip conds are gated on a scalar all-lanes predicate whose +batching rule reduces over the batch axis, so the frozen tail after the +slowest lane is genuinely skipped rather than lowered to a both-branches +select — a vmapped adaptive solve costs actual attempts, not `max_steps` +(reverse mode still stores per-slot scan residuals, so keep budgets +realistic when differentiating). + +Exactly one `SaveAt` mode is set: + +- **`SaveAt(t_1=True)`** (default): the endpoint only. +- **`SaveAt(ts=grid)`**: dense interpolation onto a fixed query grid. Output + shape is `(len(grid), ...)` regardless of how many internal steps the + controller takes; `ts` is a data leaf, so a new grid of the same length + retraces nothing. Explicit methods use cubic Hermite (4th-order accurate + between 5th-order knots); Rodas5P uses its stiff-aware fourth-order + continuous extension. Queries are observation times, not internal stops. + For an explicit constant-step solve, `exact=True` instead requires every + query to coincide with a realized knot and gathers those states directly, + with no interpolation work. +- **`SaveAt(steps=True)`**: the initial state plus accepted steps as a + contiguous prefix of a `max_steps + 1` buffer. Rejected attempts are + omitted; `fill="last"` (default) repeats the last valid row through the + tail, `fill="inf"` fills it with `inf`, and `sol.accepted` masks the valid + prefix. + +If the budget runs out before `t_1`, `sol.ok` is `False` and the outputs hold +the reached prefix. Nothing is poisoned; callers that want diverging +residuals map `jnp.where(sol.ok, x, jnp.inf)`. `sol.num_steps` counts +attempts, `sol.num_accepted` counts advances. + +### Residuals on adaptive output + +Collocation-style residuals evaluated on an adaptive +`SaveAt(steps=True)` rollout keep a static shape by evaluating the pointwise +residual on **every** padded row and zeroing the tail with the `accepted` +mask. With the default `fill="last"`, padded rows repeat the last accepted +state, so the residual there is finite wherever the endpoint is — a single +`where` is safe for both values and gradients (no double-`where` needed): + +```python +sol = jax.vmap(solve_one)(x_0s) # SaveAt(steps=True), (B, rows, ...) +rows = pointwise_residual(sol.xs, p) # evaluated on all rows in parallel +masked = jnp.where(sol.accepted, rows, 0.0) +count = jax.lax.stop_gradient(sol.accepted.sum()) # actual points, inert +residual = jnp.where( + sol.ok[:, None], masked / jnp.sqrt(count.astype(masked.dtype)), jnp.inf +).reshape(-1) # static length B * rows +``` + +Zero rows contribute nothing to a least-squares objective, so this is +exactly the fixed-shape analogue of dropping rejected steps. Two properties +to choose deliberately: normalizing by the accepted count makes the loss a +mean over actual collocation points but changes discontinuously when the +mesh changes, and lanes with more accepted steps contribute more rows — the +adaptive mesh curvature-weights the collocation. When the residual +definition needs fixed sample times and a smoother parameter dependence, +use `SaveAt(ts=fixed_grid)` instead: no mask, no mesh-dependent weighting, +and the output times carry none of the frozen-mesh ambiguity of adaptive +internal knots. + +`project` (an idempotent clamp, e.g. positivity) is applied at every point +where the field is evaluated and to every accepted state. + +## AD through adaptive stepping + +The bounded loop supports `jax.grad`, `jax.jvp`, `vmap`, and +reverse-over-forward through one adaptive solve, verified against closed +forms in `tests/test_ad.py`. + +The step-size controllers compute their error norms, decisions, and +next-step factors inside `stop_gradient`: accept/reject is a discrete branch, +and `E**(-1/order)` is singular at the exact-zero error of a flat-start +policy — without the stop-gradient and the machine-epsilon error floor the +backward pass is NaN from iteration one. + +States therefore differentiate through the solver stages on a **frozen +mesh**: the derivative holds realized step sizes and accept/reject patterns +fixed and omits the mesh-motion term `(dq/dt) * dt_k/dtheta`. Near a smooth +residual root that term is `O(‖r‖)`, so a frozen-mesh residual Jacobian is +asymptotically exact there; far from a root it can be material. Endpoint and +fixed-requested-grid outputs do not have a moving output-time axis; adaptive +`SaveAt(steps=True)` times have exactly zero tangent. + +`SaveAt(ts=...)` brackets queries with `searchsorted` — integer indices, no +gradient. Values differentiate through the bracketing states and slopes; +zero-width brackets are handled with the double-`where` trick so neither the +primal nor its transpose evaluates `0/0`. + +When several tangent or cotangent directions are needed at one primal point, +cache the linearization (`jax.linearize` / `jax.vjp`) and `vmap` the +pushforward or pullback instead of repeating the primal trajectory in a fused +`vmap(jvp)`. + +On float32 GPUs, XLA serves `dot_general` from TF32 tensor cores by default +(~1e-3 precision). When adaptive accept/reject decisions depend on +neural-network outputs, that roundoff can move the mesh; set +`jax_default_matmul_precision="highest"` in reproducibility entry points. +tinydiffeq never changes application-global JAX settings, including +`jax_enable_x64` — a float32 problem stays float32 even with x64 enabled. + +## Auxiliary outputs + +The field may return `(dx, aux)` to save quantities it already computes +without adding them to the state: + +```python +def f(x, t, args, p): + dx = -p * x + return dx, {"flow": dx, "moment": x**2} +``` + +Saved aux is a nonempty pytree of real floating arrays (leaves may differ in +dtype). It follows `SaveAt` — stored at the endpoint or accepted nodes, or +interpolated onto a requested grid with endpoint slopes obtained by JVP — and +participates in JVP/VJP. `has_aux=None` (default) detects the contract with +one `jax.eval_shape` trace; `has_aux=False` skips detection and selects the +minimal value-only path. Every saved aux leaf must remain finite: an invalid +value in a prefix-saving mode terminates at the previous accepted node, while +endpoint mode keeps the endpoint state, zero-fills aux, and sets `ok=False`. +For batched differentiation where inactive lanes may leave the model domain, +`failure_ad_reference=(x, t, p)` supplies a finite point used only for safe +linearization. + +## Stiff ODEs: Rodas5P + +`Rodas5P()` is an eight-stage, fifth-order Rosenbrock–Wanner method with an +embedded estimator and a stiff-aware fourth-order continuous extension. It is +A-stable and stiffly accurate. + +```python +import jax.numpy as jnp + +from tinydiffeq import IController, Rodas5P, solve_ode + + +def stiff_field(x, t): + # Exact solution x(t) = cos(t), with a fast transient eigenvalue -1000. + return -1000.0 * (x - jnp.cos(t)) - jnp.sin(t) + + +sol = solve_ode( + stiff_field, Rodas5P(), 0.0, 1.0, jnp.asarray(1.0), + dt_0=0.01, controller=IController(), max_steps=512, +) +``` + +At each attempted step it forms the exact JAX Jacobian $J = \partial f / +\partial x$, factors $W = I/(\gamma h) - J$ once with pivoted LU, and reuses +the factors for all eight stage solves — linear solves only, no Newton +iteration. JVP/VJP propagate through Jacobian construction and the stage +solves via the linear-solve rule $dx = A^{-1}(db - dA\,x)$; pivot selection +is not differentiated. Dense Jacobians and dense LU only — use SciML or +diffrax for sparse/Krylov stages. + +### Credit + +Rodas5P was constructed by Gerd Steinebach and introduced in +[Steinebach (2023)](https://doi.org/10.1007/s10543-023-00967-x). tinydiffeq's +JAX implementation deliberately follows SciML's authoritative MIT-licensed +implementation — the +[`OrdinaryDiffEqRosenbrock` package](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock), +its +[tableau and dense coefficients](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_tableaus.jl), +[step implementation](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_perform_step.jl), +and +[dense interpolants](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_interpolants.jl). +The tableau, stage equations, embedded estimate, and continuous extension are +ported rather than redesigned, and regression tests compare fixed steps +against SciML-produced values to prevent silent divergence. diff --git a/docs/rodas5p.md b/docs/rodas5p.md deleted file mode 100644 index 2b02514..0000000 --- a/docs/rodas5p.md +++ /dev/null @@ -1,147 +0,0 @@ -# Rodas5P: Stiff ODEs and Index-1 DAEs - -`Rodas5P()` is an eight-stage, fifth-order Rosenbrock–Wanner method with an -embedded adaptive estimator and a stiff-aware fourth-order continuous -extension. It is A-stable and stiffly accurate. In tinydiffeq it supports -ordinary ODEs and semi-explicit index-1 DAEs through the existing solve APIs. - -## Credit and implementation lineage - -Rodas5P was constructed by Gerd Steinebach and introduced and benchmarked in -[Steinebach (2023)](https://doi.org/10.1007/s10543-023-00967-x). The method was -developed in the Julia Differential Equations ecosystem, and tinydiffeq's JAX -implementation deliberately follows SciML's authoritative implementation: - -- [SciML `OrdinaryDiffEqRosenbrock` package](https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock) -- [SciML Rodas5P tableau and dense coefficients](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_tableaus.jl) -- [SciML consolidated Rosenbrock step implementation](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_perform_step.jl) -- [SciML Rosenbrock dense interpolants](https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/lib/OrdinaryDiffEqRosenbrock/src/rosenbrock_interpolants.jl) -- [SciML solver documentation and selection guidance](https://docs.sciml.ai/DiffEqDocs/stable/api/ordinarydiffeq/massmatrixdae/Rosenbrock/) - -SciML's implementation is MIT-licensed. Its source is the reference for the -port and receives explicit attribution here, in the kernel source, and in the -cross-library regression tests and benchmarks. - -The tableau, stage equations, embedded estimate, and continuous extension are -ported rather than redesigned. Regression tests compare fixed Rodas5P steps -against values produced by SciML to prevent the implementations from silently -diverging. - -## ODE use - -```python -import jax.numpy as jnp - -from tinydiffeq import IController, Rodas5P, solve_ode - - -def stiff_field(x, t): - # Exact solution x(t) = cos(t), with a fast transient eigenvalue -1000. - return -1000.0 * (x - jnp.cos(t)) - jnp.sin(t) - - -sol = solve_ode( - stiff_field, - Rodas5P(), - 0.0, - 1.0, - jnp.asarray(1.0), - dt_0=0.01, - controller=IController(), - max_steps=512, -) -``` - -At the beginning of an attempted step, Rodas5P forms the exact JAX Jacobian -(J=\partial f/\partial x), the time derivative (f_t), and - -$$ -W = \frac{I}{\gamma h} - J. -$$ - -It computes one pivoted LU factorization of (W) and reuses it for all eight -stage right-hand sides. It therefore solves linear systems but performs no -Newton or Levenberg–Marquardt iteration. - -## Semi-explicit DAE use - -For - -$$ -\dot y=f(y,z,t), \qquad 0=g(y,z,t), -$$ - -tinydiffeq constructs the flattened mass-matrix system internally: - -$$ -M\dot u=F(u,t), \qquad -u=(y,z), \quad M=\operatorname{diag}(I_y,0_z), \quad F=(f,g). -$$ - -```python -from tinydiffeq import Rodas5P, solve_semi_explicit_dae - - -sol = solve_semi_explicit_dae( - lambda y, z, t, args, p: p * z, - lambda y, z: z**2 - y - 2.0, - Rodas5P(), - 0.0, - 1.0, - jnp.asarray(1.0), - jnp.sqrt(jnp.asarray(3.0)), - p=jnp.asarray(-0.2), - dt_0=0.1, - controller=IController(), - max_steps=128, -) -``` - -`LMRootSolver` is used once to make the initial algebraic state consistent. -Every later Rodas5P stage solves - -$$ -\left(\frac{M}{\gamma h}-F_u\right)k_i=r_i -$$ - -with the reused LU factors. There is no nonlinear endpoint restoration. -Consequently, returned internal `z` values satisfy the constraint to the -method's integration accuracy, not necessarily to `LMRootSolver.atol`. - -This differs intentionally from the RK4 and Tsit5 DAE paths, which perform a -nonlinear algebraic solve at every stage and accepted endpoint. Rodas5P is -especially useful when those algebraic solves dominate runtime or when the -coupled dynamics are stiff. - -## Dense output, aux, and AD - -`SaveAt(ts=...)` uses Rodas5P's SciML/paper continuous extension for the ODE -state or combined `(y, z)` state. The three stored coefficient pytrees form a -fourth-order stiff-aware polynomial; no query-time algebraic solve occurs. - -When the differential field returns `(value, saved_aux)`, aux is evaluated and -stored at the required initial and accepted states. Requested-grid aux uses -normalized cubic Hermite interpolation with endpoint tangents obtained from -the Rodas5P polynomial and a JVP through the aux map. It is not recalculated -or root-solved at each requested time. DAE algebraic aux is separate internal -context and is never interpolated; see [Auxiliary Outputs](aux.md). - -The public derivative is the derivative of the discrete Rodas5P method. -JVP/VJP propagate through Jacobian construction, stage solves, state output, -and dense output. Each factored solve uses the mathematical linear-solve rule -$dx=A^{-1}(db-dA\,x)$, so discrete LU pivot choices are not differentiated. -As with the other adaptive methods, accept/reject decisions and step-size -selection are stop-gradiented. - -## Deliberate limits - -- Dense exact Jacobians and dense pivoted LU only; no sparse, Krylov, or - preconditioner interface. -- Constant identity or internally constructed block-diagonal mass matrices; - no public general mass-matrix API. -- Semi-explicit index-1 DAEs with a square, locally nonsingular $g_z$. -- No fully implicit residual-form DAE interface or automatic index reduction. - -For those broader capabilities, use SciML's -[`OrdinaryDiffEqRosenbrock`](https://docs.sciml.ai/DiffEqDocs/stable/api/ordinarydiffeq/massmatrixdae/Rosenbrock/) -and related DAE solvers. diff --git a/docs/sdae.md b/docs/sdae.md index fd65636..bd61f64 100644 --- a/docs/sdae.md +++ b/docs/sdae.md @@ -7,9 +7,9 @@ dy = f(y,z,t)\,dt + \sigma(y,z,t)\,dW, \qquad 0=g(y,z,t), $$ -with fixed-step Euler–Maruyama and diagonal noise in the differential state. -The algebraic Jacobian $g_z$ must be square and nonsingular along the chosen -root branch. Arrays and pytrees are supported for both `y` and `z`. +with fixed steps and diagonal noise in the differential state. The algebraic +Jacobian $g_z$ must be square and nonsingular along the chosen root branch. +Arrays and pytrees are supported for both `y` and `z`. ```python import jax @@ -54,70 +54,73 @@ sol = solve_semi_explicit_sdae( ## Algorithm and convergence -At a consistent node, the update is +`EulerMaruyama` and `SRA1` are supported. At a consistent node, +Euler–Maruyama updates $$ y_{n+1}=y_n+f(y_n,z_n,t_n)h+\sigma(y_n,z_n,t_n)\Delta W_n, $$ -followed by a root solve for -$g(y_{n+1},z_{n+1},t_{n+1})=0$. Locally writing the unique root as -$z=Z(y,t)$ shows that this is exactly Euler–Maruyama on the reduced SDE. No -Itô correction is missing: `z` is reconstructed from the constraint rather -than advanced by a separate SDE. - -The method has strong order 0.5 and weak order 1 when $g_z$ is uniformly -nonsingular near the path, the reduced drift and diffusion are globally -Lipschitz with linear growth, and the usual additional smoothness assumptions -for weak convergence hold. Additive noise gives strong order 1. A per-node -root error $\epsilon$ should be $O(\sqrt h)$ or smaller to preserve strong -order 0.5, and $O(h)$ or smaller to preserve weak order 1; normal root -tolerances are typically much tighter. +followed by a root solve for $g(y_{n+1},z_{n+1},t_{n+1})=0$. Locally writing +the unique root as $z=Z(y,t)$ shows this is exactly Euler–Maruyama on the +reduced SDE — no Itô correction is missing, since `z` is reconstructed from +the constraint rather than advanced by a separate SDE. It has strong order +0.5 and weak order 1 under the usual regularity assumptions on the reduced +system (strong order 1 for additive noise). A per-node root error +$\epsilon$ should be $O(\sqrt h)$ or smaller to preserve strong order 0.5; +normal root tolerances are far tighter. + +`SRA1` applies the Rößler additive-noise stochastic Runge–Kutta scheme to +the same reduced SDE: each step evaluates the drift at the node and at one +internal stage, with a root solve restoring consistency at the stage time +$t_n + \tfrac34 h$ and at the endpoint — two root solves per step. Its +strong order 1.5 requires the diffusion to be independent of the state; in +the SDAE setting, `z` depends on `y` through the constraint, so the +requirement is that `diffusion(y, z, t)` depends only on `t`. As in +`solve_sde`, the contract is documented, not runtime-checked, and the +endpoint diffusion evaluation reuses the current node's `(y, z, context)` +with the endpoint time. `Milstein` is deliberately not wired in: its +commutativity condition and derivative correction do not translate cleanly +through the implicit reduction $z = Z(y, t)$. ## Randomness, aux, and AD -The Brownian increments use the same sampling contract as `solve_sde`. -Arrays use the draw shape `(n_steps,) + y_0.shape`; pytrees use one flat draw -partitioned in deterministic leaf order. A fixed key defines a fixed path, -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 `LMRootSolver` configuration and -implicit-AD contract as deterministic DAEs; see -[Nonlinear-solve and AD contract](dae.md#nonlinear-solve-and-ad-contract). -nlls-gram supplies both the primal LM root and its implicit derivative, using -direct `LU()` by default for the square derivative. `gtol` and `xtol` must -remain zero, and every accepted iterate must report `CONVERGED` with Euclidean -residual norm strictly below the root `atol`. `max_steps_is_success` is retained -for source compatibility but does not make `MAX_STEPS` valid. +Noise follows the `solve_sde` sampling contract via +`solver.sample_noise(y_0, key, n_steps, dt, dtype)`: Brownian increments +for `EulerMaruyama`, an independent $(\Delta W, \Delta Z)$ pair for `SRA1`. +Arrays use the draw shape `(n_steps,) + y_0.shape`; pytrees use one flat +draw partitioned in deterministic leaf order. A fixed key defines a fixed +path, 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 `LMRootSolver` configuration and +implicit-AD contract as deterministic DAEs — see +[Nonlinear-solve and AD contract](dae.md#nonlinear-solve-and-ad-contract): +residual-only stopping, `CONVERGED` with residual norm strictly below the +root `atol`, and nlls-gram's direct square `LU()` implicit derivative. + The algebraic function may return `(residual, algebraic_aux)`; that internal -context is passed to both drift and diffusion but is not stored. The drift may -return `(drift_value, saved_aux)`, and only that saved aux becomes `sol.aux`. -Steps mode stores it at every consistent node; endpoint mode evaluates it only -at the final node. Its derivatives include both direct parameter dependence -and dependence through the implicit root. See -[Auxiliary Outputs](aux.md) for all contract variations and explicit flags. -Invalid algebraic context at initialization sets `ok=False` before any -stochastic step. In steps mode, invalid saved aux terminates at the previous -consistent node. In endpoint mode, invalid final saved aux retains the -endpoint state, returns zero aux, and sets `ok=False`. - -`sol.num_root_solves` counts logical active root calls, including the initial -consistency solve and failures, while `sol.num_root_steps` sums their LM update -counts. Both are path diagnostics with exact-zero tangents. - -Only `SaveAt(t_1=True)` and `SaveAt(steps=True)` are supported. Stochastic +context is passed to both drift and diffusion but is not stored. The drift +may return `(drift_value, saved_aux)`, and only that saved aux becomes +`sol.aux`: steps mode stores it at every consistent node, endpoint mode +evaluates it only at the final node. Its derivatives include both direct +parameter dependence and dependence through the implicit root. Invalid +algebraic context at initialization sets `ok=False` before any stochastic +step; invalid saved aux terminates at the previous consistent node (steps +mode) or zero-fills aux at the retained endpoint (endpoint mode). + +`sol.num_root_solves` counts logical active root calls — for `SRA1`, a +stage and an endpoint root per step plus the initial consistency solve — +and `sol.num_root_steps` sums their LM updates; both have exact-zero +tangents. + +Only `SaveAt(t_1=True)` and `SaveAt(steps=True)` are supported: stochastic paths are rough, so deterministic dense interpolation between nodes would be -mathematically wrong. A failed root freezes the last consistent prefix, -sets `ok=False`, and pads the remaining static buffer. Failed roots have zero -implicit tangents, and aux at a failed initial root is zero-filled, so masked -lanes can preserve successful JVPs or VJPs under `vmap`. For that contract, -pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)` at a point where the -residual, context, and saved-aux maps are finite and differentiable. The -reference is substituted before nlls only for already-inactive lanes and is -also used by inactive context, drift, diffusion, and saved-aux evaluation. A -newly attempted root must be JVP-safe at its actual initial point; the -reference cannot repair an invalid active attempt after it fails. Without one, -an all-ones best-effort default is used for inactive work. Outputs and -gradients from a failed lane are not a valid solution. +mathematically wrong. A failed root freezes the last consistent prefix, sets +`ok=False`, and pads the remaining static buffer with zero implicit +tangents, so masked lanes preserve successful JVPs/VJPs under `vmap`. For +that contract, pass `failure_ad_reference=(y_ref, z_ref, t_ref, p_ref)` at a +point where the residual, context, and saved-aux maps are finite and +differentiable — the same semantics as the +[DAE reference point](dae.md#nonlinear-solve-and-ad-contract). diff --git a/docs/sde.md b/docs/sde.md index 85c8303..b0b9c73 100644 --- a/docs/sde.md +++ b/docs/sde.md @@ -6,16 +6,18 @@ $$ dX_t = \mu(X_t, t)\,dt + \sigma(X_t, t)\,dW_t $$ -with fixed-step Euler–Maruyama and **diagonal noise**: `diffusion` returns the -same pytree structure as the state, multiplied leafwise and elementwise by -independent Brownian increments. +with fixed steps and **diagonal noise**: `diffusion` returns the same pytree +structure as the state, multiplied leafwise and elementwise by independent +Brownian increments. `n_steps` must be a static Python int; adaptive SDE +stepping requires Brownian-bridge noise that can be re-evaluated on +subdivided intervals, which is diffrax territory. ```python -from tinydiffeq import solve_sde, EulerMaruyama, SaveAt +from tinydiffeq import solve_sde, SRA1, SaveAt sol = solve_sde( - drift, diffusion, EulerMaruyama(), 0.0, 1.0, x_0, - key=jax.random.PRNGKey(0), + drift, diffusion, SRA1(), 0.0, 1.0, x_0, + key=jax.random.key(0), n_steps=256, p=(mu, sigma), save_at=SaveAt(steps=True), @@ -24,17 +26,47 @@ sol = solve_sde( `drift` and `diffusion` follow the same `(x)`, `(x, t)`, `(x, t, args)`, `(x, t, args, p)` signature convention as `solve_ode`. -The drift may return `(drift_value, saved_aux)`; diffusion remains value-only. -Aux is stored at the same endpoint or fixed grid nodes and participates in -pathwise JVP/VJP. See [Auxiliary Outputs](aux.md). -## The static-shape contract, honestly +## Solvers -`n_steps` must be a static Python int. There is no adaptive SDE stepping in -the current release — adaptivity for SDEs requires noise that can be *re-evaluated* on -subdivided intervals (a Brownian-bridge / VirtualBrownianTree construction), -so that a rejected step resamples consistently. That is diffrax territory -for now; a roadmap issue sketches what it would take here. +Three fixed-step schemes, each a frozen dataclass: + +| Solver | Strong order | Noise contract | +|---|---|---| +| `EulerMaruyama()` | 0.5 | any diagonal diffusion | +| `Milstein()` | 1.0 | diagonal *commutative*: each diffusion component depends only on its own state component | +| `SRA1()` | 1.5 | *additive*: diffusion independent of the state (may depend on time) | + +`Milstein` adds the correction $\tfrac12 g g' (\Delta W^2 - \Delta t)$, with +$g g'$ computed as the forward-mode derivative of the diffusion field in the +direction of its own value. `SRA1` is the Rößler additive-noise stochastic +Runge–Kutta scheme: two drift stages per step plus the time-Wiener integral +$I_{10}$. The contracts are documented, not runtime-checked: a +state-dependent diffusion under `SRA1` silently degrades the order, exactly +as in other SDE libraries. + +Each solver declares its own per-step randomness through +`solver.sample_noise(x_0, key, n_steps, dt, dtype)`: Brownian increments +$\Delta W$ for `EulerMaruyama`/`Milstein`, an independent pair +$(\Delta W, \Delta Z)$ for `SRA1` (from which $I_{10}/\Delta t = +\tfrac12(\Delta W + \Delta Z/\sqrt 3)$ is formed internally). + +## Explicit noise: the shocks as data + +Pass `noise=` instead of `key=` (exactly one is required) to hand `solve_sde` +the realization directly — the same pytree `sample_noise` would produce, +validated against the solver's spec: + +```python +noise = SRA1().sample_noise(x_0, key, n_steps, dt, x_0.dtype) +sol = solve_sde(drift, diffusion, SRA1(), 0.0, 1.0, x_0, noise=noise, n_steps=n_steps) +``` + +This is bit-identical to the `key=` call that would have drawn the same +`noise`, and it makes the noise **first-class data**: inspectable, storable +(e.g. as part of a training set alongside the initial condition), and +differentiable — `jax.grad` with respect to the `noise` pytree works, since +the steppers consume it as ordinary arrays. ## Key semantics: a fixed noise process @@ -47,44 +79,61 @@ d_w = jnp.sqrt(dt) * jax.random.normal(key, (n_steps,) + x_0.shape) so a fixed `key` pins the entire noise path: - **Reproducible** — the same key gives the same path, every call. -- **Differentiable with respect to `x_0` and `p`** (not `key`): with the path - held fixed, the solution map is smooth, and jvp/vjp against finite - differences are tested in `tests/test_sde.py`. This is exactly the "common - random numbers" setup simulation-based estimators want. +- **Differentiable with respect to `x_0` and `p`** (not `key`), and with + respect to `noise` when passed explicitly: with the path held fixed, the + solution map is smooth, and jvp/vjp against finite differences are tested + in `tests/test_sde.py`. This is exactly the common-random-numbers setup + simulation-based estimators want. For a pytree, tinydiffeq draws one `(n_steps, total_state_size)` array and -partitions it into leaves in JAX's deterministic leaf order. This makes the -noise identical to that of one equivalent flat array without flattening the -state during integration. Changing the tree structure changes the assignment -of random components and requires a new compilation. +partitions it into leaves in JAX's deterministic leaf order, so the noise is +identical to that of the equivalent flat array. Changing the tree structure +changes the assignment of random components and requires a new compilation. ## Orders of convergence -Euler–Maruyama is strong order 0.5 (pathwise) and weak order 1.0 (in -distribution) for multiplicative noise. The test suite verifies the strong -rate on geometric Brownian motion, where the exact solution driven by the -*same* increments is available in closed form: +The test suite verifies the strong rates on shared noise paths: because the +increments are exactly reproducible from the key, the exact solution (for +geometric Brownian motion, $X_T = X_0 \exp((\mu - \sigma^2/2)T + \sigma +W_T)$ with $W_T = \sum_k \Delta W_k$) is evaluated on the *same* path as the +numerical endpoint. Comparing against an independently sampled exact solution +would measure nothing. -$$ -X_T = X_0 \exp\left((\mu - \tfrac{\sigma^2}{2})T + \sigma W_T\right), -\qquad W_T = \textstyle\sum_k \Delta W_k . -$$ +## Auxiliary output + +Only the drift owns saved aux: -Because the increments `d_w` are exactly reproducible from the key, the test -regenerates them, computes the exact endpoint on the same path, and checks -the mean absolute error slope across `dt` levels lands in `[0.35, 0.65]`. -That shared-path construction is the right way to measure strong error — -comparing against an independently sampled exact solution measures nothing. +```python +def drift(x, t, args, p): + return drift_value, saved_aux + + +def diffusion(x, t, args, p): + return diffusion_value +``` + +Aux is stored at the endpoint or at the fixed grid nodes and is +differentiated pathwise under the fixed noise realization. `has_aux=None` +auto-detects the form; `has_aux=False` selects the no-aux scan without the +detection trace. ## SaveAt for SDEs `SaveAt(t_1=True)` (default) and `SaveAt(steps=True)` (`n_steps + 1` rows, or a padded accepted prefix after an aux failure) are supported. -`SaveAt(ts=...)` **raises**: cubic Hermite -interpolation assumes smooth trajectories and is simply wrong between the -points of a rough path. Land your grid on the step boundaries instead by -choosing `n_steps`. +`SaveAt(ts=...)` **raises**: cubic Hermite interpolation assumes smooth +trajectories and is simply wrong between the points of a rough path. Land +your grid on the step boundaries instead by choosing `n_steps`. `project` is applied to every drift/diffusion evaluation point and every accepted state, as in `solve_ode`. With `diffusion ≡ 0`, `solve_sde` reproduces `solve_ode` with `Euler()` on the same grid exactly. + +## Performance: `unroll` + +`unroll=` (a static int, default 1) unrolls that many steps per iteration of +the integration scan — identical values and gradients, fewer and larger GPU +dispatches, more compile time. For small-batch neural-network-drift +ensembles on an L40S, `unroll=8` cut reverse-mode solve time 2–3× and primal +time 1.2–1.7× (see `benchmarks/results/2026-08-04_vulcan-l40s-sde-fixed.md`). +The same argument exists on `solve_ode` for fixed stepping. diff --git a/docs/static_shapes.md b/docs/static_shapes.md deleted file mode 100644 index b83f60f..0000000 --- a/docs/static_shapes.md +++ /dev/null @@ -1,156 +0,0 @@ -# Static Shapes - -JAX jits fixed-shape programs. An adaptive integrator is naturally dynamic — -the number of steps depends on the data — so tinydiffeq offers two execution -contracts. The default `adaptive_loop="bounded"` provides exactly `max_steps` -static attempt slots in a bounded scan. Adaptive attempts are grouped into -small nested-scan chunks so a completed solve skips whole padded chunks. -`adaptive_loop="forward"` uses a dynamic `lax.while_loop` and executes only -actual attempts, while retaining the same static public output shapes and a -`max_steps` row buffer when the selected `SaveAt` mode requires one. - -Each iteration attempts one step: - -- an **accepted** attempt advances `(t, x)` and (for FSAL solvers) reuses the - last stage as the next first stage; -- a **rejected** attempt leaves the state in place and retries with a smaller - step; -- once `t_1` is reached, the remaining iterations **freeze**. - -The raw internal scan buffer contains repeated rows for rejected and frozen -iterations. That buffer is an implementation detail used by interpolation; -step output compacts it into accepted rows plus tail padding. - -The bounded loop preserves reverse-mode AD. Chunk-level and attempt-level -`lax.cond` branches keep expensive field, stage, and controller computations -out of its frozen tail. The forward loop supports primal evaluation, JVP, and -nested forward mode, but JAX cannot transpose its data-dependent while loop. -Under `vmap`, both strategies advance lanes together until the slowest lane -finishes; the forward strategy stops at that lane's actual attempt count rather -than always reaching `max_steps`. - -Fixed-step integration uses a smaller specialized scan without adaptive -controller or embedded-error work. `ConstantStepSize` accepts every attempt, -so `dt_0 = (t_1 - t_0)/n` with `max_steps = n` reproduces a fixed grid -exactly. Times are formed arithmetically as `t_0 + i * dt_0`, rather than by -repeatedly accumulating rounded steps. A small local tolerance snaps the -nominal last point to `t_1`; it is capped relative to `dt_0` and never scales -with `max_steps`. Consequently, increasing a nonbinding attempt budget cannot -stretch an earlier step or otherwise change the numerical method. - -If the budget runs out before `t_1`, `sol.ok` is `False` and the outputs hold -the reached prefix. The package never poisons values; the caller decides: - -```python -xs = jnp.where(sol.ok, sol.xs, jnp.inf) # kernels-style rejection -``` - -`sol.num_steps` counts attempts actually made, including rejections; -`sol.num_accepted` counts advances. These scalar diagnostics do not change the -fixed output shape. - -## SaveAt is the shape contract - -Exactly one of three modes: - -### `SaveAt(t_1=True)` — endpoint only (default) - -`sol.ts` is the reached time (equals `t_1` when `ok`), `sol.xs` the final -state. - -### `SaveAt(ts=grid)` — interpolation onto a fixed grid - -This is the answer to "adaptive steps vs static shapes". Internal steps -adapt freely; the output is dense interpolation onto your fixed -query grid, so each output leaf has shape -`(len(grid),) + corresponding_input_leaf.shape` **regardless of how many -steps the controller took**. Changing tolerances, initial conditions, -or curvature changes the internal knots but never the output shape — no -recompilation. A one-dimensional JAX/NumPy array or Python sequence is -accepted; times must be nondecreasing and within `[t_0, t_1]`, while repeated -times and omitted endpoints are allowed. Changing values without changing -the grid length does not recompile. - -By default these are observation times, not mandatory internal stops. The -adaptive controller chooses the same mesh regardless of the requested grid, -then the solver evaluates each point through dense interpolation. - -For an explicit ODE with `ConstantStepSize`, -`SaveAt(ts=grid, exact=True)` instead requires every query to coincide with a -realized internal endpoint and gathers the stored states directly. It avoids -interpolation and the endpoint-slope work needed by Hermite output. A -misaligned or unreached query makes `sol.ok` false. Exact mode is unavailable -for adaptive ODEs, Rodas5P, DAEs, SDEs, and SDAEs. - -The interpolation runs directly over the raw padded rows: duplicate knots -from rejections or the frozen tail form zero-width brackets, and the -bracketing `searchsorted` lands on the last duplicate at-or-before each -query, so no compaction pass is needed. Queries outside the knot span clamp -to the boundary values — in particular, when `ok` is `False`, queries beyond -the reached time return the last state (flat extrapolation) rather than -evaluating a cubic outside its bracket. - -The interpolant is 4th-order accurate between 5th-order-accurate knots: -expect grid values slightly less accurate than the knots themselves, which -is the standard dense-output trade-off. - -### `SaveAt(steps=True)` — accepted steps with padding - -`max_steps + 1` rows including the initial state. Accepted internal steps are -gathered chronologically into a contiguous prefix; rejected attempts are not -returned. `sol.accepted` is the validity mask (`accepted[0]` is always -`True`, so `accepted.sum() == num_accepted + 1`). On a successful solve, the -endpoint is at index `num_accepted`. - -- `fill="last"` (default) repeats the last valid time and state through the - padded tail. -- `fill="inf"` fills only the invalid tail with `inf`. - -If the attempt budget is exhausted, the same contract holds for the reached -accepted prefix, the tail repeats or masks its last state, and `sol.ok` is -`False`. No fake endpoint is inserted. Compaction is a stable fixed-size -gather; it performs no sorting. - -## Pytree states - -ODE, SDE, and DAE states may be arbitrary JAX pytrees, including registered -dataclasses. Each state contains at least one nonempty real floating array, -and all leaves use one dtype. Explicit and stochastic solver arithmetic maps -directly over leaves, so a single-array state retains its array execution -path. Rodas5P temporarily ravels the state because its dense Jacobian and LU -factorization couple all coordinates, then reconstructs the original pytree -at every public boundary. In every case the structure is resolved while -tracing. Changing leaf values with unchanged shapes and structure reuses a -compilation; changing the treedef or a leaf shape requires a new compilation. - -For every multi-row `SaveAt` mode the leading row dimension is added to each -leaf independently. `sol.accepted` is one shared mask for all leaves. - -## Why one compilation, precisely - -- The attempt budget `max_steps` is static; nothing else about the loop depends - on data shapes. -- Tolerances and PI coefficients (`IController(...)` / `PIController(...)`), - `dt_0`, `t_0`, `t_1`, `x_0`, - `args`, `p`, and `SaveAt.ts` are pytree **data leaves**. Only genuine - structure — the solver type, `SaveAt` mode, `fill`, `exact`, `max_steps`, - `adaptive_loop`, and the functions themselves — is static. - -An omitted tolerance or `dt_min` is represented by `None`, so switching a -jitted call between automatic and explicit values changes the controller -pytree structure and compiles once for each policy. Changing the numeric -values of already-explicit controller fields does not recompile. - -So this compiles once: - -```python -@jax.jit -def run(x_0, dt_0, controller, args): - return solve_ode(f, Tsit5(), 0.0, 1.0, x_0, args=args, dt_0=dt_0, - controller=controller, max_steps=128, - save_at=SaveAt(steps=True)) -``` - -across different curvatures (different accepted counts), tolerances, initial -steps, and initial conditions — pinned by `tests/test_recompile.py` with -`_cache_size() == 1` assertions. diff --git a/mkdocs.yml b/mkdocs.yml index ab6685e..11db43a 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -43,13 +43,10 @@ extra_javascript: nav: - Home: index.md - - Static Shapes: static_shapes.md - - Adaptive Stepping and AD: adaptive_ad.md - - Auxiliary Outputs: aux.md - - Rodas5P: rodas5p.md + - ODEs: ode.md + - SDEs: sde.md - Semi-Explicit DAEs: dae.md - Semi-Explicit SDAEs: sdae.md - - SDEs: sde.md - Markov Chains: markov_chains.md - Linear Exponential Solves: exponential.md - API Reference: api.md diff --git a/pyproject.toml b/pyproject.toml index bdbf795..1be5027 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "hatchling.build" [project] name = "tinydiffeq" -version = "2.4.0" +version = "2.5.0" description = "Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX with static shapes and composable AD" readme = "README.md" license = "MIT" diff --git a/src/tinydiffeq/__init__.py b/src/tinydiffeq/__init__.py index f16f738..c347abc 100644 --- a/src/tinydiffeq/__init__.py +++ b/src/tinydiffeq/__init__.py @@ -1,31 +1,17 @@ """Tiny differentiable ODE/SDE/DAE/SDAE solvers for JAX. solve_ode integrates dx/dt = f(x, t, args, p) with fixed-step (Euler, RK4), -adaptive explicit Tsit5, or linearly implicit Rodas5P methods. Fixed stepping -and the default adaptive path use bounded lax.scan loops with exactly max_steps -attempt slots and support forward/reverse AD, including reverse-over-forward. -Adaptive ODE/DAE solves may instead use an actual-work lax.while_loop for -primal, JVP, and nested forward AD; reverse mode is unsupported on that path. -SaveAt picks the endpoint, dense interpolation onto a fixed grid, accepted -internal steps with padding, or exact knot gathering for explicit fixed-step -ODEs. max_steps bounds attempted steps; Solution.num_steps reports actual work, -and arithmetic fixed-step times are independent of a nonbinding budget. -solve_sde is fixed-step -Euler-Maruyama with presampled -diagonal noise. solve_semi_explicit_dae handles index-1 systems with either -root-restored explicit methods or the stiff Rodas5P mass-matrix formulation, -plus differentiable saved differential-field aux and internal algebraic -context. DAESolution also reports nonlinear-root calls and LM update counts; -LMRootSolver delegates primal roots and square implicit derivatives to -nlls-gram, requires residual-only stopping, and offers previous-root and secant -stage predictors. -solve_semi_explicit_sdae applies fixed-step Euler-Maruyama to the reduced -index-1 stochastic system. solve_linear_ode applies dense, fixed-Krylov, or -adaptive matrix-free exponential actions to fixed homogeneous linear systems. -States may be arrays or pytrees of same-dtype real floating arrays. Fully -implicit solvers, general mass -matrices, full derivative-term PID control, events, continuous interpolation -objects, and adjoint methods are non-goals. +adaptive Tsit5, or linearly implicit Rodas5P methods on bounded lax.scan loops +with static shapes and composable forward/reverse AD. solve_sde integrates +diagonal-noise Ito SDEs (EulerMaruyama, Milstein, SRA1) from a PRNG key or an +explicit, differentiable noise pytree. solve_semi_explicit_dae and +solve_semi_explicit_sdae handle index-1 systems, delegating algebraic roots +and their implicit derivatives to nlls-gram. solve_linear_ode applies dense or +matrix-free Krylov exponential actions to fixed homogeneous linear systems, +and the Markov tools simulate and forecast finite-state chains. States may be +arrays or pytrees of same-dtype real floating arrays. Fully implicit solvers, +general mass matrices, events, continuous interpolation objects, and adjoint +methods are non-goals. """ from tinydiffeq.controllers import ConstantStepSize, IController, PIController @@ -58,7 +44,16 @@ from tinydiffeq.sdae import solve_semi_explicit_sdae from tinydiffeq.sde import solve_sde from tinydiffeq.solution import DAESolution, Solution -from tinydiffeq.solvers import RK4, Euler, EulerMaruyama, Rodas5P, Tsit5 +from tinydiffeq.solvers import ( + RK4, + SRA1, + Euler, + EulerMaruyama, + Milstein, + Rodas5P, + Tsit5, + diagonal_brownian_increments, +) __all__ = [ "solve_ode", @@ -73,6 +68,9 @@ "Tsit5", "Rodas5P", "EulerMaruyama", + "Milstein", + "SRA1", + "diagonal_brownian_increments", "ConstantStepSize", "IController", "PIController", diff --git a/src/tinydiffeq/_unvmap.py b/src/tinydiffeq/_unvmap.py new file mode 100644 index 0000000..c06c58e --- /dev/null +++ b/src/tinydiffeq/_unvmap.py @@ -0,0 +1,43 @@ +import jax +import numpy as np +from jax.extend import core +from jax.interpreters import ad, batching, mlir + +# A lax.cond whose predicate is batched is lowered under vmap to a select +# that executes BOTH branches for every lane, so the adaptive loops' +# skip-the-frozen-tail conds do no skipping in a vmapped solve: every attempt +# slot up to max_steps runs for every lane. unvmap_all reduces its boolean +# across any vmapped axes (its batching rule maps to an all-reduce and an +# unbatched result), so a cond gated on it keeps a scalar predicate under +# vmap and skips for real once every lane is finished. Its output is a +# nondifferentiable bool, hence the zero-JVP rule. + +unvmap_all_p = core.Primitive("unvmap_all") + + +def unvmap_all(x): + """All-reduce a boolean over any vmapped axes to a scalar predicate.""" + return unvmap_all_p.bind(x) + + +def unvmap_all_impl(x): + return jax.numpy.all(x) + + +unvmap_all_p.def_impl(unvmap_all_impl) +unvmap_all_p.def_abstract_eval(lambda x: jax.core.ShapedArray((), np.bool_)) + + +def unvmap_all_batch(args, dims): + (x,) = args + (dim,) = dims + # None is JAX's not_mapped sentinel: the reduced output has no batch axis. + return unvmap_all(jax.numpy.all(x, axis=dim)), None + + +batching.primitive_batchers[unvmap_all_p] = unvmap_all_batch +# The output is a nondifferentiable bool: no tangent contribution. +ad.defjvp(unvmap_all_p, None) +mlir.register_lowering( + unvmap_all_p, mlir.lower_fun(unvmap_all_impl, multiple_results=False) +) diff --git a/src/tinydiffeq/controllers.py b/src/tinydiffeq/controllers.py index e5914ec..2f074f2 100644 --- a/src/tinydiffeq/controllers.py +++ b/src/tinydiffeq/controllers.py @@ -65,14 +65,10 @@ class IController: Accept iff ``E = max(|err| / (atol + rtol * max(|x_0|, |x_1|))) <= 1`` (forced accept once the step reaches ``dt_min``), and propose ``dt_next = dt_used * clip(safety * E**(-1/order), factor_min, factor_max)`` - clipped to ``[dt_min, dt_max]``. This is the classic integral controller — - equal to diffrax's ``PIDController`` at its default coefficients - (p_coeff=0, d_coeff=0); there is no proportional term, hence the name. - - If omitted, ``rtol``/``atol`` default to ``1e-4``/``1e-6`` for float32 - states and ``1e-7``/``1e-9`` for float64 states. Explicit tolerances are - cast to the state dtype. ``dt_min`` defaults to ten machine epsilons in the - time dtype, scaled by ``max(1, |t_1|)``. + clipped to ``[dt_min, dt_max]``. Omitted ``rtol``/``atol`` default to + ``1e-4``/``1e-6`` for float32 states and ``1e-7``/``1e-9`` for float64; + ``dt_min`` defaults to ten machine epsilons in the time dtype, scaled by + ``max(1, |t_1|)``. """ rtol: float | None = None @@ -126,19 +122,12 @@ class PIController: """Proportional-integral step-size controller with max-norm error. In addition to the current scaled error ``E``, this controller carries the - previous accepted step's error ``E_prev`` and proposes - + previous accepted step's error ``E_prev`` (starting at one) and proposes ``dt_next = dt_used * clip(safety * E**(-(p_coeff+i_coeff)/order)`` - ``* E_prev**(p_coeff/order), factor_min, factor_max)``. - - ``E_prev`` starts at one and changes only after an accepted step. The - defaults ``p_coeff=0.4`` and ``i_coeff=0.3`` damp step-size oscillations on - harder problems; ``p_coeff=0, i_coeff=1`` reproduces :class:`IController`. - The error ratios and step-size update are stop-gradiented, while states - remain fully differentiable through the solver stages. If omitted, - ``rtol``/``atol`` default to ``1e-4``/``1e-6`` for float32 states and - ``1e-7``/``1e-9`` for float64 states. ``dt_min`` defaults to ten machine - epsilons in the time dtype, scaled by ``max(1, |t_1|)``. + ``* E_prev**(p_coeff/order), factor_min, factor_max)``. The defaults + ``p_coeff=0.4`` and ``i_coeff=0.3`` damp step-size oscillations; + ``p_coeff=0, i_coeff=1`` reproduces :class:`IController`, and the + tolerance and ``dt_min`` defaults match it. """ rtol: float | None = None diff --git a/src/tinydiffeq/dae.py b/src/tinydiffeq/dae.py index d26a952..732567d 100644 --- a/src/tinydiffeq/dae.py +++ b/src/tinydiffeq/dae.py @@ -31,6 +31,7 @@ where, zeros_like, ) +from tinydiffeq._unvmap import unvmap_all from tinydiffeq.controllers import ConstantStepSize from tinydiffeq.interpolation import ( hermite_interpolate, @@ -86,54 +87,24 @@ @jax.tree_util.register_dataclass @dataclass(frozen=True) class LMRootSolver: - """Configuration for algebraic solves in a semi-explicit DAE. + """Configuration for algebraic root solves in semi-explicit DAEs/SDAEs. The implementation is :class:`nlls_gram.LevenbergMarquardt` at its - defaults: dense ``Cholesky()`` for primal LM updates and the direct - nonsymmetric ``LU()`` implicit derivative that nlls selects automatically - for a square residual system. - - 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``. Every accepted algebraic root must have Euclidean - residual norm strictly below the root ``atol``. Root solves therefore use - residual stopping only: ``gtol`` and ``xtol`` must remain zero, and a - ``MAX_STEPS`` iterate is never treated as a differentiable root. - ``max_steps_is_success`` remains in the configuration for source - compatibility but does not broaden root acceptance. ``atol=None`` selects - ``1e-6`` in float32 and ``1e-10`` in float64. 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 + defaults. ``max_steps`` bounds one algebraic root's nonlinear iterations, + independently of the integration's time-step budget. Roots use residual + stopping only: ``gtol`` and ``xtol`` must remain zero, and every accepted + root must report ``CONVERGED`` with Euclidean residual norm strictly below + ``atol`` (``None`` selects ``1e-6`` in float32, ``1e-10`` in float64). + ``solver_options`` is 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. ``ad_solver`` remains an nlls-owned option; its default - is the direct square ``LU()`` rule. - - ``predictor="previous"`` warm-starts every explicit RK stage from the most - recent successful root. ``predictor="secant"`` instead extrapolates from - the accepted-step root through the most recent successful stage at a later - time. The secant is used only for a strictly later target; duplicate RK4 - stage times and failed stages fall back to the previous root. Predictor - values are differentiation-inert, so successful roots retain the same - implicit derivative. With multiple algebraic roots, however, a different - warm start can select a different branch; the secant mode is intended only - when the continued root is locally unique. A finite root-iteration budget - can also make predictor choice affect values or success status. + ``solver_options={"linear_solver": QR()}``; ``cache_jacobian`` and + ``geodesic_acceleration`` are fixed to ``False``. ``predictor`` selects + the explicit-stage warm start: ``"previous"`` reuses the most recent + successful root, ``"secant"`` extrapolates it to a strictly later stage + time and assumes the continued root is locally unique. """ max_steps: int = field(default=8, metadata=dict(static=True)) - max_steps_is_success: bool = field(default=False, 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)) @@ -145,8 +116,6 @@ def __post_init__(self): raise ValueError("LMRootSolver.max_steps must be a positive int") if self.max_steps <= 0: raise ValueError("LMRootSolver.max_steps must be a positive int") - if not isinstance(self.max_steps_is_success, bool): - raise TypeError("LMRootSolver.max_steps_is_success must be a bool") if self.atol is not None and self.atol <= 0: raise ValueError("LMRootSolver.atol must be positive or None") if self.gtol != 0: @@ -645,7 +614,14 @@ def skip_step(carry): return carry, out def body(carry, _): - return jax.lax.cond(carry[5] | carry[6], skip_step, attempt_step, carry) + # Scalar-predicate outer cond under vmap (see _unvmap): the frozen + # tail skips for real; the inner cond keeps per-lane freezing. + def live_lanes(carry): + return jax.lax.cond(carry[5] | carry[6], skip_step, attempt_step, carry) + + return jax.lax.cond( + unvmap_all(carry[5] | carry[6]), skip_step, live_lanes, carry + ) carry_0 = ( t_0, @@ -815,53 +791,21 @@ def solve_semi_explicit_dae( ): """Integrate a semi-explicit index-1 DAE. - The system is ``dy/dt = f(y, z, t, args, p)`` and - ``0 = g(y, z, t, args, p)``, with a square nonsingular algebraic Jacobian - ``dg/dz``. ``z_0`` is a root-finding guess: the initial algebraic state is - made consistent automatically, and its derivative is determined by the - constraint rather than by the guess. - - RK4 with fixed control, Tsit5 with fixed or adaptive control, and the - linearly implicit Rodas5P method with fixed or adaptive control are - supported. The outer ``max_steps`` bounds attempted time steps; - :class:`LMRootSolver.max_steps` separately bounds algebraic solves. RK4 - and Tsit5 restore ``g=0`` at every stage. Rodas5P uses the root solver only - for initial consistency, then advances the block mass-matrix system with - reused linear solves; later ``z`` values satisfy the constraint to the - integration accuracy rather than the root tolerance. - ``sol.num_root_solves`` counts logical active nonlinear root calls and - ``sol.num_root_steps`` sums their LM update steps. Rodas5P therefore reports - one root call regardless of its number of integration attempts. - - ``args`` is fixed data by convention. All differentiated model parameters - belong in ``p``. Initial consistency and explicit-method roots - differentiate implicitly with respect to ``(y, t, p)``; Rodas5P then - differentiates its discrete Jacobians and uses implicit derivatives for - its linear solves. - - ``f`` may return ``dy`` or ``(dy, saved_aux)``. If ``g`` returns - ``(residual, algebraic_aux)``, that second value is internal context and - ``f`` must take the full six-argument form - ``f(y, z, t, args, p, algebraic_aux)``. Only ``saved_aux`` is returned as - ``sol.aux``. It is stored at accepted nodes and interpolated on requested - deterministic grids; algebraic aux is never stored or interpolated. - ``has_aux`` and ``has_algebraic_aux`` default to abstract auto-detection; - explicit ``False`` selects the minimal paths without those traces. - - ``failure_ad_reference=(y, z, t, p)`` may provide a domain-safe point for - already-inactive ``vmap`` lanes and model aux/field evaluation. A newly - attempted root still requires its actual ``(y, z_guess, t, p)`` to be - JVP-safe; the reference does not replace an active root after it fails. - A nonfinite inexact algebraic-aux leaf at initialization prevents all - time-step work. Saved aux is checked at the initial and accepted nodes in - prefix/grid modes; endpoint mode checks it only after integration and - retains the endpoint state with zero aux if that check fails. - - ``adaptive_loop="bounded"`` keeps the reverse-mode-capable static scan. - ``adaptive_loop="forward"`` executes only actual adaptive attempts and, - under ``vmap``, stops after the slowest lane. It supports JVP and nested - forward mode but not ordinary reverse mode, matching JAX's dynamic-while - differentiation boundary. + The system is ``dy/dt = f(y, z, t, args, p)`` with + ``0 = g(y, z, t, args, p)`` and a square nonsingular ``dg/dz``. ``z_0`` + is a root guess: initial consistency is solved automatically and its + derivative comes from the constraint. RK4 and Tsit5 restore ``g = 0`` at + every stage through :class:`LMRootSolver`; Rodas5P advances the block + mass-matrix system with one reused LU factorization per attempted step, + so its later ``z`` values satisfy the constraint to integration accuracy. + Roots and their implicit derivatives are delegated to nlls-gram. ``f`` + may return ``(dy, saved_aux)``; ``g`` may return + ``(residual, algebraic_aux)``, in which case ``f`` takes + ``(y, z, t, args, p, algebraic_aux)``. + ``failure_ad_reference=(y, z, t, p)`` provides a domain-safe point for + inactive ``vmap`` lanes. ``adaptive_loop`` follows + :func:`tinydiffeq.solve_ode`. Returns a :class:`DAESolution` with + root-solve diagnostics. """ if dt_0 is None: raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)") @@ -1592,8 +1536,12 @@ def skip_step(carry): return carry, out def body(carry, _): - terminated = carry[8] | carry[9] - return jax.lax.cond(terminated, skip_step, attempt_step, carry) + def live_lanes(carry): + return jax.lax.cond(carry[8] | carry[9], skip_step, attempt_step, carry) + + return jax.lax.cond( + unvmap_all(carry[8] | carry[9]), skip_step, live_lanes, carry + ) carry_0 = ( t_0, diff --git a/src/tinydiffeq/exponential.py b/src/tinydiffeq/exponential.py index 15de6dd..ac13829 100644 --- a/src/tinydiffeq/exponential.py +++ b/src/tinydiffeq/exponential.py @@ -420,18 +420,13 @@ def solve_linear_ode(operator, method, t_0, t_1, x_0, *, save_at=None): """Solve ``dx/dt = A(x)`` for a fixed homogeneous linear operator. ``operator`` is either a square matrix using the column convention - ``A @ x`` or a callable that maps the state pytree to an identically + ``A @ x`` or a callable mapping the state pytree to an identically structured pytree. ``DenseExponential`` materializes a callable operator - with forward-mode Jacobian columns before applying a dense matrix - exponential. ``KrylovExponential`` and ``AdaptiveKrylovExponential`` only - evaluate operator actions and are therefore suitable for structured - matrix-free pytrees. - - Endpoint output is the default. ``SaveAt(ts=...)`` evaluates independent - exponential actions at the requested times. The solve supports ordinary - JVPs and VJPs through the initial state and through differentiable arrays - used by the operator. The operator must be autonomous, homogeneous, and - linear; affine and nonlinear exponential methods are a separate problem. + before a dense matrix exponential; the Krylov methods evaluate only + operator actions. Endpoint output is the default, and ``SaveAt(ts=...)`` + evaluates independent exponential actions at the requested times. JVPs + and VJPs flow through the initial state and differentiable operator + arrays. The operator must be autonomous, homogeneous, and linear. """ if not isinstance(method, _EXPONENTIAL_METHODS): raise TypeError( diff --git a/src/tinydiffeq/ode.py b/src/tinydiffeq/ode.py index 64045bc..72839ef 100644 --- a/src/tinydiffeq/ode.py +++ b/src/tinydiffeq/ode.py @@ -25,6 +25,7 @@ zero_tangent, zeros_like, ) +from tinydiffeq._unvmap import unvmap_all from tinydiffeq.controllers import ConstantStepSize from tinydiffeq.interpolation import ( hermite_interpolate, @@ -132,47 +133,25 @@ def solve_ode( has_aux=None, failure_ad_reference=None, adaptive_loop="bounded", + unroll=1, ): """Integrate ``dx/dt = f(x, t, args, p)`` from ``t_0`` to ``t_1 > t_0``. - The vector field may be declared ``f(x)``, ``f(x, t)``, ``f(x, t, args)``, - or ``f(x, t, args, p)`` — always in that order. ``x`` is a pytree whose - nonempty leaves share one real floating dtype. ``args`` is pass-through - data that is by convention not an AD target, and ``p`` holds differentiable - parameters (any pytree); - jvp/vjp with respect to ``p`` and ``x_0`` are first-class. - - Fixed stepping and the default ``adaptive_loop="bounded"`` use bounded - ``lax.scan`` loops with exactly ``max_steps`` attempt slots, so shapes are - static and curvature-dependent step counts never retrace. Bounded adaptive - attempts are grouped into static chunks so one ``lax.cond`` skips an entire - padded chunk after completion. ``adaptive_loop="forward"`` instead uses a - true ``lax.while_loop`` and stops after the actual attempts. It supports - primal evaluation, JVP, and nested forward-mode AD; ordinary reverse mode - is unavailable because JAX cannot transpose a dynamic ``while_loop``. - Under ``vmap``, the forward loop advances lanes together until the slowest - live trajectory finishes rather than running every lane to ``max_steps``. - ``dt_0`` is required (no auto-initial-step heuristic). - Each attempt is clipped to the remaining horizon; the clipped step also - feeds the controller's next-step proposal, which doubles as the growth - guard — near-flat fields otherwise grow steps into quarter-horizon leaps. - ``project`` (e.g. a positivity clamp, assumed idempotent) is applied at - every point where ``f`` is evaluated and to every accepted state. Returns - a :class:`Solution`; ``sol.ok`` is False if the budget ran out before - ``t_1`` or a required saved output was invalid (outputs remain a finite - reached prefix or endpoint, never poisoned). - - The field may return either ``dx`` or ``(dx, aux)``. ``has_aux=None`` - detects the form with an abstract trace; ``has_aux=False`` selects the - minimal no-aux path without that trace. Saved aux is a nonempty pytree of - real floating arrays. It follows ``SaveAt`` and participates in JVP/VJP. - By default, requested-grid aux uses cubic Hermite interpolation with - endpoint slopes obtained by JVP, including for Rodas5P's dense state path. - With ``SaveAt(ts=..., exact=True)``, an explicit fixed-step solve instead - selects realized knots and evaluates aux directly at those requested knots. - - The time dtype follows the state dtype; the library never - sets ``jax_enable_x64`` — do that in your application. + The field may be declared ``f(x)``, ``f(x, t)``, ``f(x, t, args)``, or + ``f(x, t, args, p)``; ``x`` is an array or pytree with one real floating + dtype, ``args`` is inert data, and ``p`` holds differentiable parameters. + ``dt_0`` is required. Fixed stepping and the default + ``adaptive_loop="bounded"`` use bounded ``lax.scan`` loops with exactly + ``max_steps`` attempt slots and support forward and reverse AD; + ``adaptive_loop="forward"`` runs a dynamic ``lax.while_loop`` (primal, + JVP, and nested forward mode only). ``project`` (an idempotent clamp) is + applied at every field evaluation and accepted state. The field may + return ``(dx, aux)``; saved aux follows ``SaveAt`` and participates in + AD. ``unroll`` (a static int, fixed stepping only) unrolls that many + steps per iteration of the integration scan — identical values, + fewer/larger GPU dispatches, more compile time. Returns a + :class:`Solution`; ``sol.ok`` reports whether ``t_1`` was reached with + every requested output valid — outputs are never poisoned. """ if dt_0 is None: raise ValueError("dt_0 is required (tinydiffeq has no initial-step heuristic)") @@ -196,6 +175,10 @@ def solve_ode( raise ValueError( 'adaptive_loop="forward" requires an adaptive error controller' ) + if not isinstance(unroll, int) or isinstance(unroll, bool) or unroll < 1: + raise ValueError("unroll must be a static int of at least 1") + if unroll != 1 and controller.uses_error_estimate: + raise ValueError("unroll requires fixed stepping (ConstantStepSize)") f = canonicalize_field(f) is_rodas = isinstance(solver, Rodas5P) is_fixed = isinstance(controller, ConstantStepSize) @@ -445,7 +428,15 @@ def skip_step(carry): return carry, out def body(carry, _): - return jax.lax.cond(carry[6] | carry[7], skip_step, attempt_step, carry) + # The outer cond keeps a scalar predicate under vmap (see _unvmap), + # so the frozen tail is skipped for real once every lane finishes; + # the inner cond preserves per-lane freezing while any lane is live. + def live_lanes(carry): + return jax.lax.cond(carry[6] | carry[7], skip_step, attempt_step, carry) + + return jax.lax.cond( + unvmap_all(carry[6] | carry[7]), skip_step, live_lanes, carry + ) def fixed_attempt_step(carry): t, x, f_cur, done, num_accepted = carry @@ -485,7 +476,10 @@ def fixed_skip_step(carry): return carry, output def fixed_body(carry, _): - return jax.lax.cond(carry[3], fixed_skip_step, fixed_attempt_step, carry) + def live_lanes(carry): + return jax.lax.cond(carry[3], fixed_skip_step, fixed_attempt_step, carry) + + return jax.lax.cond(unvmap_all(carry[3]), fixed_skip_step, live_lanes, carry) def uniform_fixed_body(carry, step_index): t, x, f_cur, _, num_accepted = carry @@ -550,7 +544,7 @@ def skip_chunk(chunk_carry, chunk_valid): return chunk_carry, repeat_output(output) def outer(chunk_carry, chunk_valid): - inactive = chunk_carry[6] | chunk_carry[7] + inactive = unvmap_all(chunk_carry[6] | chunk_carry[7]) return jax.lax.cond( inactive, skip_chunk, run_chunk, chunk_carry, chunk_valid ) @@ -589,11 +583,11 @@ def outer(chunk_carry, chunk_valid): if static_uniform_horizon: step_indices = jnp.arange(max_steps, dtype=jnp.int32) fixed_final, rows = jax.lax.scan( - uniform_fixed_body, fixed_carry_0, step_indices + uniform_fixed_body, fixed_carry_0, step_indices, unroll=unroll ) else: fixed_final, rows = jax.lax.scan( - fixed_body, fixed_carry_0, None, length=max_steps + fixed_body, fixed_carry_0, None, length=max_steps, unroll=unroll ) t_final, x_final, _, done, num_accepted = fixed_final num_steps = num_accepted @@ -635,7 +629,9 @@ def outer(chunk_carry, chunk_valid): _, ) = final_carry else: - final_carry, rows = jax.lax.scan(body, carry_0, None, length=max_steps) + final_carry, rows = jax.lax.scan( + body, carry_0, None, length=max_steps, unroll=unroll + ) ( t_final, x_final, diff --git a/src/tinydiffeq/save_at.py b/src/tinydiffeq/save_at.py index 722665a..a200fe7 100644 --- a/src/tinydiffeq/save_at.py +++ b/src/tinydiffeq/save_at.py @@ -7,27 +7,17 @@ @jax.tree_util.register_dataclass @dataclass(frozen=True) class SaveAt: - """What ``solve_ode``/``solve_sde`` return. Exactly one mode must be set. + """What the solve functions return. Exactly one mode must be set. - - ``t_1=True``: the endpoint only (the default in the solve functions). - - ``ts=grid``: dense interpolation of the internal steps onto a fixed, - sorted query grid in ``[t_0, t_1]``. Explicit methods use cubic Hermite; - Rodas5P uses its stiff-aware fourth-order extension. Output shape is - ``(len(ts), ...)`` regardless of how many internal steps the controller - takes, so changing curvature never changes shapes or recompiles. - ``ts`` is a data leaf; a different grid of the same length retraces - nothing. ODE, deterministic DAE, and linear exponential solves only. - For an explicit constant-step ODE, ``exact=True`` instead requires every - query to coincide with an internal step endpoint. It gathers those states - directly and avoids both interpolation and its extra endpoint-slope field - evaluation. - - ``steps=True``: the initial state and accepted internal steps as a - chronological prefix of a ``max_steps + 1`` buffer. Rejected attempts - are omitted. ``fill="last"`` (default) pads the tail with the final - valid row; ``fill="inf"`` pads the tail with ``inf``. The returned - ``Solution.accepted`` mask distinguishes the valid prefix from padding. - - ``fill`` only applies to ``steps=True``. + ``t_1=True`` (the solver default) returns the endpoint only. ``ts=grid`` + interpolates the internal steps onto a fixed query grid — output shape is + ``(len(ts), ...)`` however many steps the controller takes, and ``ts`` is + a data leaf; with ``exact=True`` an explicit constant-step ODE instead + gathers states at queries that must coincide with realized knots. + ``steps=True`` returns the initial state and accepted steps as the valid + prefix of a ``max_steps + 1`` buffer, padded with the last valid row + (``fill="last"``) or ``inf`` (``fill="inf"``) and masked by + ``Solution.accepted``. """ t_1: bool = field(default=False, metadata=dict(static=True)) diff --git a/src/tinydiffeq/sdae.py b/src/tinydiffeq/sdae.py index d510c97..9e19e43 100644 --- a/src/tinydiffeq/sdae.py +++ b/src/tinydiffeq/sdae.py @@ -1,4 +1,4 @@ -"""Fixed-step Euler--Maruyama for semi-explicit index-1 SDAEs.""" +"""Fixed-step stochastic integration for semi-explicit index-1 SDAEs.""" import jax import jax.numpy as jnp @@ -20,6 +20,7 @@ take, where, ) +from tinydiffeq._unvmap import unvmap_all from tinydiffeq.dae import ( LMRootSolver, _canonicalize_cached_dae_field, @@ -29,9 +30,8 @@ _prepare_failure_ad_reference, ) from tinydiffeq.save_at import SaveAt -from tinydiffeq.sde import _diagonal_brownian_increments from tinydiffeq.solution import DAESolution -from tinydiffeq.solvers import EulerMaruyama +from tinydiffeq.solvers import INV_SQRT_3, SRA1, EulerMaruyama def solve_semi_explicit_sdae( @@ -56,41 +56,23 @@ def solve_semi_explicit_sdae( ): """Integrate a semi-explicit index-1 Ito SDAE with diagonal noise. - The system is ``dy = drift(y, z, t) dt + diffusion(y, z, t) dW`` and - ``0 = g(y, z, t)``. Euler--Maruyama advances the differential state on a - fixed uniform grid, then an algebraic root solve restores consistency at - every node. This is Euler--Maruyama applied to the reduced SDE obtained - from the locally unique root ``z = Z(y, t)``. - ``sol.num_root_solves`` counts logical active nonlinear root calls and - ``sol.num_root_steps`` sums their LM update steps. - - ``drift`` may return ``value`` or ``(value, saved_aux)``. If ``g`` returns - ``(residual, algebraic_aux)``, that internal context is passed to both - ``drift(y, z, t, args, p, algebraic_aux)`` and the corresponding - six-argument ``diffusion``. Only drift-owned ``saved_aux`` is exposed as - ``sol.aux`` and it is stored at consistent stochastic nodes; stochastic - interpolation is deliberately unsupported. ``has_aux`` and - ``has_algebraic_aux`` default to abstract auto-detection, while explicit - ``False`` selects the minimal no-aux paths. - - A fixed key defines - one common-random-number path for JVP/VJP with respect to ``y_0`` and - ``p``. ``z_0`` is only a root guess and has zero tangent by contract. - ``failure_ad_reference=(y, z, t, p)`` may provide a domain-safe point for - already-inactive ``vmap`` lanes and model aux/field evaluation. A newly - attempted root still requires its actual ``(y, z_guess, t, p)`` to be - JVP-safe; the reference does not replace an active root after it fails. - A nonfinite inexact algebraic-aux leaf at initialization prevents all - stochastic time-step work. Saved aux is checked at every node in steps - mode; endpoint mode checks it only after integration and retains the - endpoint state with zero aux if that check fails. + The system is ``dy = drift(y, z, t) dt + diffusion(y, z, t) dW`` with + ``0 = g(y, z, t)``. ``solver`` is ``EulerMaruyama`` or ``SRA1``, applied + to the reduced SDE obtained from the locally unique root ``z = Z(y, t)``: + the differential state advances on a fixed uniform grid of ``n_steps`` + steps, and a root solve restores consistency at every node and at SRA1's + drift stage. SRA1's strong order 1.5 requires a diffusion that depends + only on time. A fixed ``key`` defines one common-random-numbers path; + JVP/VJP with respect to ``y_0`` and ``p`` are pathwise, and ``z_0`` is a + root guess with zero tangent. Aux contracts, ``failure_ad_reference``, + and failure behavior follow ``solve_semi_explicit_dae``. """ if not isinstance(n_steps, int) or isinstance(n_steps, bool): raise TypeError("n_steps must be a static Python int") if n_steps < 1: raise ValueError("n_steps must be at least 1") - if not isinstance(solver, EulerMaruyama): - raise TypeError("semi-explicit SDAEs currently support EulerMaruyama") + if not isinstance(solver, (EulerMaruyama, SRA1)): + raise TypeError("semi-explicit SDAEs support EulerMaruyama and SRA1") if save_at is None: save_at = SaveAt(t_1=True) if save_at.ts is not None: @@ -112,7 +94,8 @@ def solve_semi_explicit_sdae( ) dt = (t_1 - t_0) / n_steps time_grid = jnp.linspace(t_0, t_1, n_steps + 1) - d_w = _diagonal_brownian_increments(y_0, key, n_steps, dt, time_dtype) + noise = solver.sample_noise(y_0, key, n_steps, dt, time_dtype) + is_sra1 = isinstance(solver, SRA1) raw_drift = drift raw_diffusion = diffusion @@ -243,7 +226,7 @@ def attempt_step(carry, inputs): num_root_solves, num_root_steps, ) = carry - t_step, t_next, d_w_step = inputs + t_step, t_next, noise_step = inputs active = ~failed y_ref, z_ref, t_ref, p_ref = failure_ad_reference y_eval = where(active, y, y_ref) @@ -261,14 +244,79 @@ def attempt_step(carry, inputs): diffusion_output(y_eval, z_eval, t_eval, p_eval, context_eval), "diffusion(y, z, t)", ) - y_candidate = add_scaled( - y, - (dt, drift_value), - (1.0, multiply(diffusion_value, d_w_step)), - ) - z_candidate, root_ok, root_solves, root_steps = solve_root( - y_candidate, t_next, z, active - ) + if is_sra1: + d_w_step, d_z_step = noise_step + chi = jax.tree.map( + lambda w, v: 0.5 * (w + v * INV_SQRT_3), d_w_step, d_z_step + ) + # Additive-noise contract: the diffusion may depend only on time, + # so its endpoint evaluation reuses the node's (y, z, context). + t_next_eval = jnp.where(active, t_next, t_ref) + diffusion_next = checked_value( + diffusion_output(y_eval, z_eval, t_next_eval, p_eval, context_eval), + "diffusion(y, z, t)", + ) + t_stage = t_step + 0.75 * dt + y_stage = add_scaled( + y, + (0.75 * dt, drift_value), + (1.5, multiply(diffusion_next, chi)), + ) + z_stage, stage_root_ok, stage_solves, stage_steps = solve_root( + y_stage, t_stage, z, active + ) + if has_algebraic_aux: + context_stage, stage_context_ok = evaluate_context( + y_stage, z_stage, t_stage, stage_root_ok + ) + stage_ok = stage_root_ok & stage_context_ok + else: + context_stage = None + stage_ok = stage_root_ok + y_stage_eval = where(stage_ok, y_stage, y_ref) + z_stage_eval = where(stage_ok, z_stage, z_ref) + t_stage_eval = jnp.where(stage_ok, t_stage, t_ref) + context_stage_eval = ( + where(stage_ok, context_stage, context_reference) + if has_algebraic_aux + else None + ) + stage_drift_raw, _ = split_field_output( + drift_output( + y_stage_eval, + z_stage_eval, + t_stage_eval, + p_eval, + context_stage_eval, + ), + has_aux, + ) + stage_drift_value = checked_value(stage_drift_raw, "drift(y, z, t)") + y_candidate = add_scaled( + y, + (dt / 3.0, drift_value), + (2.0 * dt / 3.0, stage_drift_value), + (1.0, multiply(diffusion_next, d_w_step)), + ( + 1.0, + multiply(add_scaled(diffusion_value, (-1.0, diffusion_next)), chi), + ), + ) + z_candidate, endpoint_ok, endpoint_solves, endpoint_steps = solve_root( + y_candidate, t_next, z_stage, stage_ok + ) + root_ok = stage_ok & endpoint_ok + root_solves = stage_solves + endpoint_solves + root_steps = stage_steps + endpoint_steps + else: + y_candidate = add_scaled( + y, + (dt, drift_value), + (1.0, multiply(diffusion_value, noise_step)), + ) + z_candidate, root_ok, root_solves, root_steps = solve_root( + y_candidate, t_next, z, active + ) if has_algebraic_aux: context_candidate, context_ok = evaluate_context( y_candidate, z_candidate, t_next, root_ok & active @@ -337,10 +385,18 @@ def skip_step(carry, _): return carry, out def body(carry, inputs): + def live_lanes(pair): + return jax.lax.cond( + pair[0][5], + lambda pair: skip_step(*pair), + lambda pair: attempt_step(*pair), + pair, + ) + return jax.lax.cond( - carry[5], + unvmap_all(carry[5]), lambda pair: skip_step(*pair), - lambda pair: attempt_step(*pair), + live_lanes, (carry, inputs), ) @@ -373,7 +429,7 @@ def body(carry, inputs): ) = jax.lax.scan( body, carry_0, - (time_grid[:-1], time_grid[1:], d_w), + (time_grid[:-1], time_grid[1:], noise), ) ok = ~failed & (num_accepted == n_steps) if save_at.t_1: diff --git a/src/tinydiffeq/sde.py b/src/tinydiffeq/sde.py index 781be75..1ba278d 100644 --- a/src/tinydiffeq/sde.py +++ b/src/tinydiffeq/sde.py @@ -15,32 +15,12 @@ take, where, ) +from tinydiffeq._unvmap import unvmap_all from tinydiffeq.ode import canonicalize_field, identity_project from tinydiffeq.save_at import SaveAt from tinydiffeq.solution import Solution -def _diagonal_brownian_increments(x_0, key, n_steps, dt, dtype): - """Generate one diagonal-noise draw for an array or pytree state.""" - leaves, treedef = jax.tree.flatten(x_0) - if treedef == jax.tree.structure(0): - return jnp.sqrt(dt) * jax.random.normal( - key, (n_steps,) + x_0.shape, dtype=dtype - ) - sizes = [leaf.size for leaf in leaves] - flat_noise = jnp.sqrt(dt) * jax.random.normal( - key, (n_steps, sum(sizes)), dtype=dtype - ) - noise_leaves = [] - start = 0 - for leaf, size in zip(leaves, sizes, strict=True): - noise_leaves.append( - flat_noise[:, start : start + size].reshape((n_steps,) + leaf.shape) - ) - start += size - return jax.tree.unflatten(treedef, noise_leaves) - - def solve_sde( drift, diffusion, @@ -49,43 +29,42 @@ def solve_sde( t_1, x_0, *, - key, + key=None, n_steps, + noise=None, p=None, args=None, save_at=None, project=None, has_aux=None, failure_ad_reference=None, + unroll=1, ): - """Integrate the Ito SDE ``dx = drift dt + diffusion d_w`` (diagonal noise) - on the fixed grid of ``n_steps`` uniform steps from ``t_0`` to ``t_1 > t_0``. + """Integrate the Ito SDE ``dx = drift dt + diffusion d_w`` with diagonal + noise on a fixed grid of ``n_steps`` uniform steps from ``t_0`` to + ``t_1 > t_0``. ``drift`` and ``diffusion`` follow the same signature convention as - ``solve_ode`` — ``(x)``, ``(x, t)``, ``(x, t, args)``, or - ``(x, t, args, p)``. ``n_steps`` must be a static Python int (the honest - static-shape contract; there is currently no adaptive SDE stepping). The - Brownian increments are presampled from ``key``. Arrays retain the exact - ``(n_steps,) + x_0.shape`` random draw. Pytree states use one shared flat - draw, partitioned into leaves in JAX's deterministic pytree leaf order. - Thus a fixed key gives a fixed noise process: reproducible across calls - and differentiable with respect to ``x_0`` and ``p`` (not ``key``). - - ``SaveAt(ts=...)`` raises — cubic Hermite interpolation is wrong for - rough paths; use ``t_1`` (default) or ``steps`` (here ``n_steps + 1`` - rows unless a saved-aux failure terminates the accepted prefix). - - ``drift`` may return either its value or ``(value, aux)``. The optional - real-floating aux pytree is stored at the same fixed nodes as ``xs`` and - is differentiated pathwise under the fixed random key. ``diffusion`` is - value-only. ``has_aux=None`` auto-detects the drift form; - ``has_aux=False`` avoids the abstract detection trace and selects the - original no-aux scan. + ``solve_ode``. ``solver`` is ``EulerMaruyama``, ``Milstein``, or ``SRA1``, + each declaring its per-step noise through + ``solver.sample_noise(x_0, key, n_steps, dt, dtype)``. Exactly one of + ``key`` and ``noise`` must be provided: a fixed ``key`` presamples a + fixed, reproducible noise process, differentiable with respect to ``x_0`` + and ``p``; an explicit ``noise`` pytree (validated against the solver's + spec) is additionally differentiable as data. ``SaveAt(ts=...)`` raises — + interpolation is wrong for rough paths. ``drift`` may return + ``(value, aux)``; ``diffusion`` is value-only. ``unroll`` (a static int) + unrolls that many steps per iteration of the underlying ``lax.scan`` — + identical values, fewer/larger GPU dispatches, more compile time. """ if not isinstance(n_steps, int): raise TypeError("n_steps must be a static Python int") if n_steps < 1: raise ValueError("n_steps must be at least 1") + if not isinstance(unroll, int) or isinstance(unroll, bool) or unroll < 1: + raise ValueError("unroll must be a static int of at least 1") + if (key is None) == (noise is None): + raise ValueError("solve_sde requires exactly one of key or noise") if save_at is None: save_at = SaveAt(t_1=True) if save_at.ts is not None: @@ -102,7 +81,36 @@ def solve_sde( t_0 = jnp.asarray(t_0, time_dtype) t_1 = jnp.asarray(t_1, time_dtype) dt = (t_1 - t_0) / n_steps - d_w = _diagonal_brownian_increments(x_0, key, n_steps, dt, time_dtype) + if noise is None: + noise = solver.sample_noise(x_0, key, n_steps, dt, time_dtype) + else: + # sample_noise only reads shapes from x_0, so the abstract trace is + # the solver's authoritative noise spec. + reference = jax.eval_shape( + lambda noise_key: solver.sample_noise( + x_0, noise_key, n_steps, dt, time_dtype + ), + jax.random.key(0), + ) + noise = jax.tree.map(jnp.asarray, noise) + if jax.tree.structure(noise) != jax.tree.structure(reference): + raise ValueError( + "noise must match the pytree structure of " + "solver.sample_noise(x_0, key, n_steps, dt, dtype)" + ) + for leaf, ref in zip( + jax.tree.leaves(noise), jax.tree.leaves(reference), strict=True + ): + if leaf.shape != ref.shape: + raise ValueError( + f"noise leaf shape {leaf.shape} does not match the " + f"solver's expected {ref.shape}" + ) + if leaf.dtype != ref.dtype: + raise TypeError( + f"noise leaf dtype {leaf.dtype} must match the state " + f"dtype {ref.dtype}" + ) time_grid = jnp.linspace(t_0, t_1, n_steps + 1) def project_state(x): @@ -142,12 +150,14 @@ def g_diffusion(x, t): return value def body(x, inputs): - t, d_w_step = inputs - x_1 = solver.step(g_drift, g_diffusion, t, x, dt, d_w_step, project_state) + t, noise_step = inputs + x_1 = solver.step(g_drift, g_diffusion, t, x, dt, noise_step, project_state) return x_1, x_1 if save_at.steps else None if save_at.t_1 or not has_aux: - x_final, step_states = jax.lax.scan(body, x_0, (time_grid[:-1], d_w)) + x_final, step_states = jax.lax.scan( + body, x_0, (time_grid[:-1], noise), unroll=unroll + ) num_accepted = jnp.asarray(n_steps, jnp.int32) num_steps = num_accepted ok = jnp.asarray(True) @@ -194,14 +204,14 @@ def auxiliary(inputs): def aux_attempt(carry, inputs): x, aux, t, failed, count, num_steps = carry - t_step, t_next, d_w_step = inputs + t_step, t_next, noise_step = inputs x_candidate = solver.step( g_drift, g_diffusion, t_step, x, dt, - d_w_step, + noise_step, project_state, ) aux_candidate, aux_ok = evaluate_aux( @@ -230,10 +240,20 @@ def aux_skip(carry, inputs): return carry, (t, x, aux, jnp.asarray(False)) def aux_body(carry, inputs): + # Scalar-predicate outer cond under vmap (see _unvmap): once every + # lane has failed, the frozen tail skips for real. + def live_lanes(pair): + return jax.lax.cond( + pair[0][3], + lambda pair: aux_skip(*pair), + lambda pair: aux_attempt(*pair), + pair, + ) + return jax.lax.cond( - carry[3], + unvmap_all(carry[3]), lambda pair: aux_skip(*pair), - lambda pair: aux_attempt(*pair), + live_lanes, (carry, inputs), ) @@ -255,7 +275,9 @@ def aux_body(carry, inputs): num_steps, ), rows, - ) = jax.lax.scan(aux_body, carry_0, (time_grid[:-1], time_grid[1:], d_w)) + ) = jax.lax.scan( + aux_body, carry_0, (time_grid[:-1], time_grid[1:], noise), unroll=unroll + ) ts_s, xs_s, aux_s, advance_s = rows all_times = jnp.concatenate([t_0[None], ts_s]) all_states = prepend(x_0, xs_s) diff --git a/src/tinydiffeq/solution.py b/src/tinydiffeq/solution.py index 4b61c7d..7a7a6dd 100644 --- a/src/tinydiffeq/solution.py +++ b/src/tinydiffeq/solution.py @@ -9,24 +9,14 @@ class Solution: """Result of ``solve_ode``/``solve_sde``/``solve_linear_ode``. - - ``ts``/``xs``: times and states in the shape dictated by ``SaveAt`` — - scalar/endpoint for ``t_1``; for multi-row modes, every state pytree leaf - receives a leading ``len(ts)`` or ``max_steps + 1`` axis. - - ``ok``: scalar bool, True iff the integration reached ``t_1`` within the - attempt budget and every required saved output was valid. The package - never poisons outputs; callers that want diverging residuals map - ``jnp.where(sol.ok, x, jnp.inf)`` over leaves. - - ``num_accepted``: number of accepted steps (excluding the initial - state). - - ``num_steps``: scalar integer array counting logical attempted steps, - including rejected attempts. Public solver outputs always populate it; - the optional default only preserves direct-construction compatibility. - - ``accepted``: ``steps`` mode only (otherwise None): validity mask for - the contiguous accepted-step prefix. Row 0 (the initial state) is - always True, so ``accepted.sum() == num_accepted + 1``. - - ``aux``: optional floating pytree returned by the differential field - (or stochastic drift), with the same leading saved-time axis as - ``xs``. It is ``None`` when the field has no auxiliary output. + ``ts``/``xs`` hold times and states in the shape dictated by ``SaveAt``. + ``ok`` is a scalar bool: the integration reached ``t_1`` and every + required saved output was valid. Outputs are never poisoned; callers that + want diverging values map ``jnp.where(sol.ok, x, jnp.inf)`` over leaves. + ``num_accepted`` counts accepted steps, ``num_steps`` counts logical + attempts including rejections, ``accepted`` masks the valid prefix in + ``steps`` mode (row 0 is always True), and ``aux`` holds the field's + saved auxiliary pytree with the same leading saved-time axis as ``xs``. """ ts: jax.Array @@ -41,26 +31,16 @@ class Solution: @jax.tree_util.register_dataclass @dataclass(frozen=True) class DAESolution: - """Result of the deterministic or stochastic semi-explicit DAE solver. - - ``ts`` and ``ys`` follow the same :class:`tinydiffeq.SaveAt` shape - contract as :class:`Solution`. ``zs`` holds algebraic states and ``aux`` - holds the optional differential-field auxiliary-output pytree. An - algebraic auxiliary output is internal context passed into the - differential field; it is not stored directly. RK4/Tsit5 - internal-step and endpoint values are evaluated at converged roots. - Rodas5P instead satisfies the algebraic equation to integration accuracy - after its initial consistency root. Requested values are dense - interpolants and need not satisfy the constraint exactly. ``ok`` is true - only when initialization, all required stage and saved-output operations - succeeded, and the integration reached ``t_1``. ``num_root_solves`` counts - logical active algebraic root calls, including failed calls and the initial - consistency solve. ``num_root_steps`` sums their nonlinear LM update steps. - ``num_steps`` counts logical attempted integration steps, including rejected - attempts, with the same semantics as :class:`Solution`. - All three are path diagnostics with exact-zero tangents. Under ``vmap``, masked - lanes may still execute physically even though they do not increment these - logical per-lane counters. + """Result of the deterministic or stochastic semi-explicit DAE solvers. + + ``ts``/``ys``/``zs`` follow the :class:`Solution` shape contract, with + ``zs`` the algebraic states. Explicit-method saved values sit at + converged roots; Rodas5P satisfies the constraint to integration accuracy + after its initial consistency root, and requested-grid interpolants need + not satisfy it exactly. ``num_root_solves`` counts logical active root + calls (including failures and the initial consistency solve) and + ``num_root_steps`` sums their LM update steps; like ``num_steps``, both + are path diagnostics with exact-zero tangents. """ ts: jax.Array diff --git a/src/tinydiffeq/solvers.py b/src/tinydiffeq/solvers.py index dd83f8a..f1cc9f7 100644 --- a/src/tinydiffeq/solvers.py +++ b/src/tinydiffeq/solvers.py @@ -1,6 +1,7 @@ from dataclasses import dataclass import jax +import jax.numpy as jnp from tinydiffeq._tree import add_scaled, multiply, weighted_sum @@ -185,14 +186,13 @@ def step_fixed(self, g, t, x, dt, f_0, project): class Rodas5P: """Fifth-order Rodas5P Rosenbrock--Wanner method. - Rodas5P is an eight-stage, linearly implicit method with an embedded - error estimate and a stiff-aware fourth-order continuous extension. It is - supported by :func:`tinydiffeq.solve_ode` and - :func:`tinydiffeq.solve_semi_explicit_dae`; both use one dense LU - factorization per attempted step and reuse it across all stages. - - The implementation follows Steinebach (2023) and SciML's - ``OrdinaryDiffEqRosenbrock.Rodas5P`` implementation: + An eight-stage, linearly implicit method with an embedded error estimate + and a stiff-aware fourth-order continuous extension, supported by + :func:`tinydiffeq.solve_ode` and + :func:`tinydiffeq.solve_semi_explicit_dae` with one dense LU + factorization reused across the stages of each attempted step. The + implementation follows Steinebach (2023) and SciML's + ``OrdinaryDiffEqRosenbrock.Rodas5P``: - https://doi.org/10.1007/s10543-023-00967-x - https://github.com/SciML/OrdinaryDiffEq.jl/tree/master/lib/OrdinaryDiffEqRosenbrock @@ -203,14 +203,137 @@ class Rodas5P: has_error_estimate = True +def diagonal_brownian_increments(x_0, key, n_steps, dt, dtype): + """Draw ``n_steps`` diagonal Brownian increments ``sqrt(dt) * N(0, 1)``. + + Arrays retain the exact ``(n_steps,) + x_0.shape`` draw. Pytree states use + one shared flat draw, partitioned into leaves in JAX's deterministic + pytree leaf order. + """ + leaves, treedef = jax.tree.flatten(x_0) + if treedef == jax.tree.structure(0): + return jnp.sqrt(dt) * jax.random.normal( + key, (n_steps,) + x_0.shape, dtype=dtype + ) + sizes = [leaf.size for leaf in leaves] + flat_noise = jnp.sqrt(dt) * jax.random.normal( + key, (n_steps, sum(sizes)), dtype=dtype + ) + noise_leaves = [] + start = 0 + for leaf, size in zip(leaves, sizes, strict=True): + noise_leaves.append( + flat_noise[:, start : start + size].reshape((n_steps,) + leaf.shape) + ) + start += size + return jax.tree.unflatten(treedef, noise_leaves) + + +# SDE steppers share the contract +# `step(g_drift, g_diffusion, t, x, dt, noise, project) -> x_1` where `noise` +# is one per-step slice of the pytree produced by the solver's own +# `sample_noise(x_0, key, n_steps, dt, dtype)`, so explicit noise handed to +# `solve_sde` is validated against exactly what the solver expects. + + @jax.tree_util.register_dataclass @dataclass(frozen=True) class EulerMaruyama: - """Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only.""" + """Euler-Maruyama for Ito SDEs with diagonal noise. Fixed-step only. + + Strong order 0.5 for multiplicative noise. ``sample_noise`` returns the + Brownian increments with the same pytree structure as the state and a + leading ``n_steps`` axis. + """ + + order = 1 + strong_order = 0.5 + + def sample_noise(self, x_0, key, n_steps, dt, dtype): + return diagonal_brownian_increments(x_0, key, n_steps, dt, dtype) + + def step(self, g_drift, g_diffusion, t, x, dt, noise, project): + return project( + add_scaled( + x, (dt, g_drift(x, t)), (1.0, multiply(g_diffusion(x, t), noise)) + ) + ) + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class Milstein: + """Milstein for Ito SDEs with diagonal noise. Fixed-step only. + + Strong order 1.0 under the diagonal commutativity condition: each + diffusion component may depend only on its own state component. The + correction ``(1/2) g g' (d_w^2 - dt)`` evaluates ``g g'`` as the + forward-mode derivative of the diffusion field in the direction of its + own value, which equals the diagonal term exactly in that case. + ``sample_noise`` matches ``EulerMaruyama``. + """ order = 1 + strong_order = 1.0 + + def sample_noise(self, x_0, key, n_steps, dt, dtype): + return diagonal_brownian_increments(x_0, key, n_steps, dt, dtype) - def step(self, g_drift, g_diffusion, t, x, dt, d_w, project): + def step(self, g_drift, g_diffusion, t, x, dt, noise, project): + g_value = g_diffusion(x, t) + _, dg_g = jax.jvp(lambda state: g_diffusion(state, t), (x,), (g_value,)) + correction = jax.tree.map(lambda dg, w: 0.5 * dg * (w * w - dt), dg_g, noise) return project( - add_scaled(x, (dt, g_drift(x, t)), (1.0, multiply(g_diffusion(x, t), d_w))) + add_scaled( + x, + (dt, g_drift(x, t)), + (1.0, multiply(g_value, noise)), + (1.0, correction), + ) + ) + + +INV_SQRT_3 = 3.0**-0.5 + + +@jax.tree_util.register_dataclass +@dataclass(frozen=True) +class SRA1: + """Rossler SRA1 stochastic Runge-Kutta for Ito SDEs with additive + diagonal noise. Fixed-step only. + + Strong order 1.5 when the diffusion is independent of the state (it may + depend on time). ``sample_noise`` returns ``(d_w, d_z)``: two independent + ``sqrt(dt) * N(0, 1)`` draws per step. The time-Wiener integral + ``I_10 / dt`` is realized internally as ``(d_w + d_z / sqrt(3)) / 2``, + reproducing its variance ``dt^3 / 3`` and covariance ``dt^2 / 2`` with + the increment. + """ + + order = 2 + strong_order = 1.5 + + def sample_noise(self, x_0, key, n_steps, dt, dtype): + key_w, key_z = jax.random.split(key) + return ( + diagonal_brownian_increments(x_0, key_w, n_steps, dt, dtype), + diagonal_brownian_increments(x_0, key_z, n_steps, dt, dtype), + ) + + def step(self, g_drift, g_diffusion, t, x, dt, noise, project): + d_w, d_z = noise + g_0 = g_diffusion(x, t) + g_1 = g_diffusion(x, t + dt) + chi = jax.tree.map(lambda w, z: 0.5 * (w + z * INV_SQRT_3), d_w, d_z) + k_1 = g_drift(x, t) + stage = add_scaled(x, (0.75 * dt, k_1), (1.5, multiply(g_1, chi))) + k_2 = g_drift(stage, t + 0.75 * dt) + return project( + add_scaled( + x, + (dt / 3.0, k_1), + (2.0 * dt / 3.0, k_2), + (1.0, multiply(g_1, d_w)), + (1.0, multiply(add_scaled(g_0, (-1.0, g_1)), chi)), + ) ) diff --git a/tests/test_aux.py b/tests/test_aux.py index aee7ef5..6e1edeb 100644 --- a/tests/test_aux.py +++ b/tests/test_aux.py @@ -308,7 +308,7 @@ def differential(y, z, t, args, p, algebraic_aux): max_steps=1, save_at=save_at, has_aux=True, - root_solver=LMRootSolver(max_steps_is_success=False), + root_solver=LMRootSolver(), failure_ad_reference=(1.0, 1.0, 0.0, 0.0), ) diff --git a/tests/test_dae.py b/tests/test_dae.py index 8337797..21ced4b 100644 --- a/tests/test_dae.py +++ b/tests/test_dae.py @@ -46,7 +46,6 @@ def constraint(y, z, t, args, p): return z - y defaults = _build_algebraic_solver(constraint, LMRootSolver(), False) - assert not LMRootSolver().max_steps_is_success # Everything algorithmic is nlls-gram's default; only the two invariants # this package owns are pinned. assert isinstance(defaults.linear_solver, Cholesky) @@ -109,7 +108,7 @@ def test_lm_root_solver_options_normalize_and_reject_fixed_keys(): def test_max_steps_policy_requires_root_residual_and_strict_batched_derivative(): - def one_lane(z_0, p, max_steps_is_success): + def one_lane(z_0, p): return solve_semi_explicit_dae( lambda y, z: jnp.zeros_like(y), lambda y, z, t, args, p: z - p, @@ -121,26 +120,20 @@ def one_lane(z_0, p, max_steps_is_success): p=p, dt_0=0.1, max_steps=1, - root_solver=LMRootSolver( - max_steps=1, - max_steps_is_success=max_steps_is_success, - atol=1e-12, - ), + root_solver=LMRootSolver(max_steps=1, atol=1e-12), ) p = jnp.asarray(1.0) - forgiving = one_lane(jnp.asarray(-1e6), p, True) - strict = one_lane(jnp.asarray(0.0), p, False) - assert not bool(forgiving.ok) + starved = one_lane(jnp.asarray(-1e6), p) + strict = one_lane(jnp.asarray(0.0), p) + assert not bool(starved.ok) assert not bool(strict.ok) - assert int(forgiving.num_root_solves) == 1 - assert int(forgiving.num_root_steps) == 1 - assert forgiving.zs == jnp.asarray(-1e6) + assert int(starved.num_root_solves) == 1 + assert int(starved.num_root_steps) == 1 + assert starved.zs == jnp.asarray(-1e6) def strict_batch(parameter): - return jax.vmap(lambda z_0: one_lane(z_0, parameter, False))( - jnp.asarray([1.0, 0.0]) - ) + return jax.vmap(lambda z_0: one_lane(z_0, parameter))(jnp.asarray([1.0, 0.0])) def endpoint(parameter): result = strict_batch(parameter) @@ -162,7 +155,7 @@ def test_dae_roots_reject_nonresidual_stopping_rules(): LMRootSolver(xtol=1e-6) -def test_compat_max_steps_success_does_not_broaden_nlls_root_ad(): +def test_root_budget_exhaustion_is_never_a_differentiable_root(): guesses = jnp.asarray([1.0, 0.0]) def solve_one(parameter, guess): @@ -177,11 +170,7 @@ def solve_one(parameter, guess): p=parameter, dt_0=0.1, max_steps=1, - root_solver=LMRootSolver( - max_steps=1, - max_steps_is_success=True, - atol=1e-12, - ), + root_solver=LMRootSolver(max_steps=1, atol=1e-12), ) def endpoints(parameters): @@ -803,7 +792,7 @@ def test_initial_root_failure_and_time_budget_failure(): jnp.asarray(0.0), dt_0=0.1, max_steps=10, - root_solver=LMRootSolver(max_steps_is_success=False), + root_solver=LMRootSolver(), ) assert not bool(failed_root.ok) assert int(failed_root.num_accepted) == 0 @@ -842,7 +831,6 @@ def test_adaptive_stage_root_failure_retries_with_smaller_step(predictor): controller=IController(), root_solver=LMRootSolver( max_steps=1, - max_steps_is_success=False, atol=1e-8, predictor=predictor, ), @@ -872,7 +860,7 @@ def one_lane(y, z, p): p=p, dt_0=0.1, max_steps=1, - root_solver=LMRootSolver(max_steps_is_success=False), + root_solver=LMRootSolver(), failure_ad_reference=(1.0, 1.0, 0.0, 0.0), ) @@ -949,8 +937,6 @@ def test_validation(): LMRootSolver(max_steps=0) with pytest.raises(ValueError, match="atol must be positive or None"): LMRootSolver(atol=0.0) - with pytest.raises(TypeError, match="max_steps_is_success must be a bool"): - LMRootSolver(max_steps_is_success=1) with pytest.raises(ValueError, match="gtol must be zero"): LMRootSolver(gtol=-1.0) with pytest.raises(ValueError, match="xtol must be zero"): diff --git a/tests/test_float64_subprocess.py b/tests/test_float64_subprocess.py index 186cdac..1006a8d 100644 --- a/tests/test_float64_subprocess.py +++ b/tests/test_float64_subprocess.py @@ -431,6 +431,60 @@ def continuous_forecast(value): """) +def test_sra1_with_x64_disabled_stays_float32_and_differentiable(): + run_script(r""" +import jax +import jax.numpy as jnp + +from tinydiffeq import SRA1, solve_sde + + +assert not jax.config.x64_enabled +n = 32 +x_0 = jnp.asarray(1.0) +key = jax.random.key(0) +noise = SRA1().sample_noise(x_0, key, n, jnp.asarray(1.0 / n), x_0.dtype) + + +def endpoint(x_0, noise): + return solve_sde( + lambda x: -0.5 * x, + lambda x: 0.1 * jnp.ones_like(x), + SRA1(), + 0.0, + 1.0, + x_0, + noise=noise, + n_steps=n, + ).xs + + +jaxpr = jax.make_jaxpr(endpoint)(x_0, noise) +assert "f64" not in str(jaxpr), jaxpr + +value = jax.jit(endpoint)(x_0, noise) +keyed = solve_sde( + lambda x: -0.5 * x, + lambda x: 0.1 * jnp.ones_like(x), + SRA1(), + 0.0, + 1.0, + x_0, + key=key, + n_steps=n, +).xs +assert value.dtype == jnp.float32 +assert jnp.array_equal(value, keyed) + +grad_x, grad_noise = jax.grad(endpoint, argnums=(0, 1))(x_0, noise) +assert grad_x.dtype == jnp.float32 +assert jnp.isfinite(grad_x) +for leaf in jax.tree.leaves(grad_noise): + assert leaf.dtype == jnp.float32 + assert jnp.all(jnp.isfinite(leaf)) +""") + + def test_pytree_states_preserve_float32_and_float64(): run_script(r""" import jax diff --git a/tests/test_gpu.py b/tests/test_gpu.py index e8a36f0..74a55f1 100644 --- a/tests/test_gpu.py +++ b/tests/test_gpu.py @@ -3,15 +3,19 @@ import pytest from tinydiffeq import ( + RK4, + SRA1, AdaptiveKrylovExponential, AssociativeMarkov, ContinuousTimeMarkovChain, DenseExponential, DiscreteMarkovChain, + Euler, EulerMaruyama, IController, KrylovExponential, MatrixFreeContinuousTimeMarkovChain, + Milstein, Rodas5P, SaveAt, Tsit5, @@ -625,6 +629,192 @@ def endpoint(p): assert bool(jnp.isfinite(leaf)) +def test_float32_fixed_euler_rk4_ode_on_gpu(): + gpu = gpu_devices()[0] + n = 128 + + @jax.jit + def run(x_0): + euler = solve_ode( + lambda x: -F32_DECAY * x, Euler(), 0.0, 1.0, x_0, dt_0=1.0 / n, max_steps=n + ).xs + rk4 = solve_ode( + lambda x: -F32_DECAY * x, RK4(), 0.0, 1.0, x_0, dt_0=1.0 / n, max_steps=n + ).xs + return euler, rk4 + + with jax.default_device(gpu): + euler, rk4 = jax.block_until_ready(run(jnp.asarray(1.0, jnp.float32))) + + exact = float(jnp.exp(-F32_DECAY)) + for out, tol in ((euler, 1e-3), (rk4, 1e-5)): + assert out.dtype == jnp.float32 + assert out.devices().pop().platform == "gpu" + assert abs(float(out) - exact) < tol + + +@pytest.mark.parametrize( + ("solver", "diffusion"), + [ + (Milstein(), lambda x: 0.1 * x), + (SRA1(), lambda x: 0.1 * jnp.ones_like(x)), + ], + ids=["milstein", "sra1"], +) +def test_float32_milstein_and_sra1_run_on_gpu(solver, diffusion): + gpu = gpu_devices()[0] + + @jax.jit + def run(x_0): + return solve_sde( + lambda x: -F32_DECAY * x, + diffusion, + solver, + 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 + + +# The kernels use case: vmap over B trajectories with per-trajectory x_0 and +# explicit noise, jit the whole batch, and take reverse-mode gradients of a +# scalar residual with respect to (x_0, p, noise). n_steps covers 31 and 127 +# so both a scan XLA may unroll and one it will not take the same paths. + + +def sra1_batch_noise(batch, n_steps, dtype, seed): + keys = jax.random.split(jax.random.key(seed), batch) + dt = jnp.asarray(1.0 / n_steps, dtype) + return jax.vmap( + lambda k: SRA1().sample_noise(jnp.zeros((), dtype), k, n_steps, dt, dtype) + )(keys) + + +@pytest.mark.parametrize("n_steps", [31, 127]) +def test_float32_vmapped_sra1_ensemble_matches_per_trajectory_on_gpu(n_steps): + gpu = gpu_devices()[0] + dtype = jnp.float32 + batch = 8 + x_0s = jnp.linspace(0.5, 2.0, batch, dtype=dtype) + noise = sra1_batch_noise(batch, n_steps, dtype, seed=3) + + def one(x_0, w, z): + return solve_sde( + lambda x: -F32_DECAY * x, + lambda x: jnp.asarray(0.1, x.dtype) * jnp.ones_like(x), + SRA1(), + 0.0, + 1.0, + x_0, + noise=(w, z), + n_steps=n_steps, + ).xs + + with jax.default_device(gpu): + batched = jax.block_until_ready(jax.jit(jax.vmap(one))(x_0s, *noise)) + looped = jnp.stack( + [one(x_0s[i], noise[0][i], noise[1][i]) for i in range(batch)] + ) + + assert batched.dtype == dtype + assert batched.devices().pop().platform == "gpu" + assert jnp.allclose(batched, looped, atol=1e-5) + + +@pytest.mark.parametrize("n_steps", [31, 127]) +def test_float32_grad_of_vmapped_sra1_ensemble_on_gpu(n_steps): + gpu = gpu_devices()[0] + dtype = jnp.float32 + batch = 8 + x_0s = jnp.linspace(0.5, 2.0, batch, dtype=dtype) + theta = jnp.asarray(0.7, dtype) + noise = sra1_batch_noise(batch, n_steps, dtype, seed=11) + + def endpoint(x_0, p, w, z): + return solve_sde( + lambda x, t, args, p: -p * x, + lambda x, t, args, p: jnp.asarray(0.1, x.dtype) * jnp.ones_like(x), + SRA1(), + 0.0, + 1.0, + x_0, + p=p, + noise=(w, z), + n_steps=n_steps, + ).xs + + def loss(x_0s, p, noise): + endpoints = jax.vmap(lambda x_0, w, z: endpoint(x_0, p, w, z))(x_0s, *noise) + return jnp.sum((endpoints - 1.0) ** 2) + + with jax.default_device(gpu): + value, (grad_x, grad_p, grad_noise) = jax.block_until_ready( + jax.jit(jax.value_and_grad(loss, argnums=(0, 1, 2)))(x_0s, theta, noise) + ) + + for leaf in jax.tree.leaves((value, grad_x, grad_p, grad_noise)): + assert leaf.dtype == dtype + assert leaf.devices().pop().platform == "gpu" + assert bool(jnp.all(jnp.isfinite(leaf))) + + def one_loss(x_0, w, z): + return (endpoint(x_0, theta, w, z) - 1.0) ** 2 + + per_trajectory = jnp.stack( + [jax.grad(one_loss)(x_0s[i], noise[0][i], noise[1][i]) for i in range(batch)] + ) + assert jnp.allclose(grad_x, per_trajectory, atol=1e-4) + + +@pytest.mark.parametrize("n_steps", [31, 127]) +def test_float32_grad_of_vmapped_rk4_ensemble_on_gpu(n_steps): + gpu = gpu_devices()[0] + dtype = jnp.float32 + batch = 8 + x_0s = jnp.linspace(0.5, 2.0, batch, dtype=dtype) + theta = jnp.asarray(0.7, dtype) + + def endpoint(x_0, p): + return solve_ode( + lambda x, t, args, p: -p * x, + RK4(), + 0.0, + 1.0, + x_0, + p=p, + dt_0=1.0 / n_steps, + max_steps=n_steps, + ).xs + + def loss(x_0s, p): + endpoints = jax.vmap(lambda x_0: endpoint(x_0, p))(x_0s) + return jnp.sum((endpoints - 1.0) ** 2) + + with jax.default_device(gpu): + grad_x, grad_p = jax.block_until_ready( + jax.jit(jax.grad(loss, argnums=(0, 1)))(x_0s, theta) + ) + + decay = jnp.exp(-theta) + expected_x = 2.0 * (x_0s * decay - 1.0) * decay + expected_p = jnp.sum(-2.0 * (x_0s * decay - 1.0) * x_0s * decay) + for leaf in (grad_x, grad_p): + assert leaf.dtype == dtype + assert leaf.devices().pop().platform == "gpu" + assert jnp.allclose(grad_x, expected_x, atol=1e-3) + assert jnp.allclose(grad_p, expected_p, atol=1e-2) + + 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 diff --git a/tests/test_sdae.py b/tests/test_sdae.py index ed822e9..e545e7d 100644 --- a/tests/test_sdae.py +++ b/tests/test_sdae.py @@ -4,8 +4,10 @@ import pytest from tinydiffeq import ( + SRA1, EulerMaruyama, LMRootSolver, + Milstein, SaveAt, solve_sde, solve_semi_explicit_sdae, @@ -26,7 +28,7 @@ def constraint(y, z, t, args, p): return z - y, {"scaled": p["scale"] * z, "square": z**2} -def sdae(key, n_steps, *, save_at=None, y_0=Y_0, p=None): +def sdae(key, n_steps, *, save_at=None, y_0=Y_0, p=None, solver=None): if p is None: p = { "mu": jnp.asarray(MU), @@ -37,7 +39,7 @@ def sdae(key, n_steps, *, save_at=None, y_0=Y_0, p=None): drift, diffusion, constraint, - EulerMaruyama(), + EulerMaruyama() if solver is None else solver, 0.0, T, jnp.asarray(y_0), @@ -120,6 +122,83 @@ def exact(key): assert 0.35 < slope < 0.7, (slope, errors) +def additive_sigma(t, p): + # Time-dependent but state-independent, so SRA1's additive contract holds + # for the reduced SDE and its endpoint diffusion evaluation is exercised. + return p["sigma"] * (1.0 + 0.5 * t) + + +def test_sra1_sdae_matches_reduced_sde_on_identical_noise_path(): + key = jax.random.key(12) + n = 32 + p = {"mu": jnp.asarray(MU), "sigma": jnp.asarray(SIGMA)} + full = solve_semi_explicit_sdae( + lambda y, z, t, args, p: p["mu"] * z, + lambda y, z, t, args, p: additive_sigma(t, p) * jnp.ones_like(y), + lambda y, z: z - y, + SRA1(), + 0.0, + T, + jnp.asarray(Y_0), + jnp.asarray(0.7), + key=key, + n_steps=n, + p=p, + save_at=SaveAt(steps=True), + ) + reduced = solve_sde( + lambda y, t, args, p: p["mu"] * y, + lambda y, t, args, p: additive_sigma(t, p) * jnp.ones_like(y), + SRA1(), + 0.0, + T, + jnp.asarray(Y_0), + key=key, + n_steps=n, + p=p, + save_at=SaveAt(steps=True), + ) + assert bool(full.ok) + assert int(full.num_steps) == n + # One initial consistency root, then a stage and an endpoint root per step. + assert int(full.num_root_solves) == 2 * n + 1 + assert jnp.allclose(full.ys, reduced.xs, atol=1e-7, rtol=1e-7) + assert jnp.allclose(full.zs, full.ys, atol=1e-7) + + +def test_sra1_sdae_aux_and_state_jvp_vjp_under_fixed_key(): + key = jax.random.key(21) + + def output(mu): + p = {"mu": mu, "sigma": jnp.asarray(SIGMA), "scale": mu} + sol = sdae(key, 64, p=p, save_at=SaveAt(steps=True), solver=SRA1()) + return jnp.sum(sol.ys + sol.zs + sol.aux["scaled"]) + + mu = jnp.asarray(MU) + tangent = jax.jvp(output, (mu,), (jnp.ones_like(mu),))[1] + cotangent = jax.grad(output)(mu) + eps = 1e-6 + finite_difference = (output(mu + eps) - output(mu - eps)) / (2 * eps) + assert jnp.abs(tangent - finite_difference) < 2e-5 + assert jnp.abs(cotangent - finite_difference) < 2e-5 + + +def test_sdae_rejects_milstein(): + with pytest.raises(TypeError, match="EulerMaruyama and SRA1"): + solve_semi_explicit_sdae( + lambda y, z: z, + lambda y, z: 0.0 * y, + lambda y, z: z - y, + Milstein(), + 0.0, + 1.0, + jnp.asarray(1.0), + jnp.asarray(1.0), + key=jax.random.key(0), + n_steps=4, + ) + + def test_sdae_root_guess_has_zero_tangent_and_ts_raises(): key = jax.random.key(1) @@ -203,7 +282,7 @@ def stochastic(y, z, t, args, p, algebraic_aux): save_at=SaveAt(steps=True), has_aux=True, has_algebraic_aux=True, - root_solver=LMRootSolver(max_steps_is_success=False), + root_solver=LMRootSolver(), ) assert not bool(sol.ok) assert int(sol.accepted.sum()) == int(sol.num_accepted) + 1 @@ -244,7 +323,7 @@ def zero_diffusion(y, z, t, args, p, algebraic_aux): save_at=save_at, has_aux=True, has_algebraic_aux=True, - root_solver=LMRootSolver(max_steps_is_success=False), + root_solver=LMRootSolver(), failure_ad_reference=(1.0, 1.0, 0.0, 0.0), ) diff --git a/tests/test_sde.py b/tests/test_sde.py index f32c455..6aa5af8 100644 --- a/tests/test_sde.py +++ b/tests/test_sde.py @@ -3,7 +3,15 @@ import numpy as np import pytest -from tinydiffeq import Euler, EulerMaruyama, SaveAt, solve_ode, solve_sde +from tinydiffeq import ( + SRA1, + Euler, + EulerMaruyama, + Milstein, + SaveAt, + solve_ode, + solve_sde, +) # Geometric Brownian motion dX = mu X dt + sigma X d_w has the exact solution # X_T = X_0 exp((mu - sigma^2/2) T + sigma W_T). solve_sde presamples its @@ -202,3 +210,398 @@ def test_steps_mode_shapes_and_flags(): assert int(sol.num_steps) == n assert bool(jnp.all(sol.accepted)) assert sol.ts[0] == 0.0 and sol.ts[-1] == T + + +# Additive-noise Ornstein-Uhlenbeck dX = -theta X dt + sigma d_w for the SRA1 +# tests: SRA1's strong order 1.5 holds only for state-independent diffusion. + +THETA, OU_SIGMA = 1.0, 0.5 + + +def ou_drift(x, t, args, p): + return -THETA * x + + +def ou_diffusion(x, t, args, p): + return OU_SIGMA * jnp.ones_like(x) + + +@pytest.mark.parametrize("solver", [EulerMaruyama(), Milstein(), SRA1()]) +@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64]) +def test_explicit_noise_matches_key(solver, dtype): + key = jax.random.PRNGKey(11) + n = 32 + x_0 = jnp.asarray(X_0, dtype) + noise = solver.sample_noise(x_0, key, n, jnp.asarray(T / n, dtype), dtype) + a = solve_sde( + ou_drift, + ou_diffusion, + solver, + 0.0, + T, + x_0, + key=key, + n_steps=n, + save_at=SaveAt(steps=True), + ) + b = solve_sde( + ou_drift, + ou_diffusion, + solver, + 0.0, + T, + x_0, + noise=noise, + n_steps=n, + save_at=SaveAt(steps=True), + ) + assert a.xs.dtype == dtype + assert a.ts.dtype == dtype + assert jnp.array_equal(a.xs, b.xs) + assert jnp.array_equal(a.ts, b.ts) + + +def test_key_noise_exclusivity_raises(): + x_0 = jnp.asarray(X_0) + noise = EulerMaruyama().sample_noise( + x_0, jax.random.PRNGKey(0), 8, jnp.asarray(T / 8), x_0.dtype + ) + with pytest.raises(ValueError, match="exactly one"): + solve_sde( + drift, diffusion, EulerMaruyama(), 0.0, T, x_0, n_steps=8, p=(MU, SIGMA) + ) + with pytest.raises(ValueError, match="exactly one"): + solve_sde( + drift, + diffusion, + EulerMaruyama(), + 0.0, + T, + x_0, + key=jax.random.PRNGKey(0), + noise=noise, + n_steps=8, + p=(MU, SIGMA), + ) + + +def test_explicit_noise_wrong_shape_raises(): + with pytest.raises(ValueError, match="shape"): + solve_sde( + drift, + diffusion, + EulerMaruyama(), + 0.0, + T, + jnp.asarray(X_0), + noise=jnp.zeros((7,)), + n_steps=8, + p=(MU, SIGMA), + ) + + +def test_explicit_noise_wrong_structure_raises(): + # SRA1 expects the (d_w, d_z) pair, not a bare increment array. + with pytest.raises(ValueError, match="structure"): + solve_sde( + ou_drift, + ou_diffusion, + SRA1(), + 0.0, + T, + jnp.asarray(X_0), + noise=jnp.zeros((8,)), + n_steps=8, + ) + + +def test_explicit_noise_wrong_dtype_raises(): + with pytest.raises(TypeError, match="dtype"): + solve_sde( + drift, + diffusion, + EulerMaruyama(), + 0.0, + T, + jnp.asarray(X_0), + noise=jnp.zeros((8,), dtype=jnp.float32), + n_steps=8, + p=(MU, SIGMA), + ) + + +def test_grad_wrt_explicit_noise_vs_finite_differences(): + n = 16 + x_0 = jnp.asarray(X_0) + noise = EulerMaruyama().sample_noise( + x_0, jax.random.PRNGKey(5), n, jnp.asarray(T / n), x_0.dtype + ) + + def endpoint(noise): + return solve_sde( + drift, + diffusion, + EulerMaruyama(), + 0.0, + T, + x_0, + noise=noise, + n_steps=n, + p=(MU, SIGMA), + ).xs + + grad = jax.grad(endpoint)(noise) + eps = 1e-6 + for i in (0, n // 2, n - 1): + bump = jnp.zeros_like(noise).at[i].set(eps) + fd = (endpoint(noise + bump) - endpoint(noise - bump)) / (2 * eps) + assert jnp.abs(grad[i] - fd) < 1e-6, i + + +def test_grad_wrt_sra1_noise_pair_vs_finite_differences(): + n = 16 + x_0 = jnp.asarray(X_0) + noise = SRA1().sample_noise( + x_0, jax.random.PRNGKey(6), n, jnp.asarray(T / n), x_0.dtype + ) + + def endpoint(noise): + return solve_sde( + ou_drift, ou_diffusion, SRA1(), 0.0, T, x_0, noise=noise, n_steps=n + ).xs + + grad_w, grad_z = jax.grad(endpoint)(noise) + eps = 1e-6 + for leaf, grad_leaf in ((0, grad_w), (1, grad_z)): + for i in (0, n - 1): + bump = jnp.zeros((n,)).at[i].set(eps) + if leaf == 0: + plus = (noise[0] + bump, noise[1]) + minus = (noise[0] - bump, noise[1]) + else: + plus = (noise[0], noise[1] + bump) + minus = (noise[0], noise[1] - bump) + fd = (endpoint(plus) - endpoint(minus)) / (2 * eps) + assert jnp.abs(grad_leaf[i] - fd) < 1e-6, (leaf, i) + + +@pytest.mark.parametrize("dtype", [jnp.float32, jnp.float64]) +def test_vmap_explicit_noise_matches_loop(dtype): + n = 16 + x_0s = jnp.asarray([0.5, 1.0, 2.0], dtype) + keys = jax.random.split(jax.random.PRNGKey(9), 3) + d_w, d_z = jax.vmap( + lambda k: SRA1().sample_noise( + jnp.asarray(X_0, dtype), k, n, jnp.asarray(T / n, dtype), dtype + ) + )(keys) + + def one(x_0, w, z): + return solve_sde( + ou_drift, ou_diffusion, SRA1(), 0.0, T, x_0, noise=(w, z), n_steps=n + ).xs + + batched = jax.vmap(one)(x_0s, d_w, d_z) + looped = jnp.stack([one(x_0s[i], d_w[i], d_z[i]) for i in range(3)]) + assert batched.dtype == dtype + assert jnp.max(jnp.abs(batched - looped)) < 200 * jnp.finfo(dtype).eps + + +def test_unroll_matches_rolled_values_and_gradients(): + key = jax.random.PRNGKey(13) + n = 32 + x_0 = jnp.asarray(X_0) + + def endpoint(x_0, unroll): + return solve_sde( + ou_drift, + ou_diffusion, + SRA1(), + 0.0, + T, + x_0, + key=key, + n_steps=n, + unroll=unroll, + ).xs + + rolled = solve_sde( + ou_drift, + ou_diffusion, + SRA1(), + 0.0, + T, + x_0, + key=key, + n_steps=n, + save_at=SaveAt(steps=True), + unroll=1, + ) + unrolled = solve_sde( + ou_drift, + ou_diffusion, + SRA1(), + 0.0, + T, + x_0, + key=key, + n_steps=n, + save_at=SaveAt(steps=True), + unroll=4, + ) + assert jnp.array_equal(rolled.xs, unrolled.xs) + grad_rolled = jax.grad(lambda x: endpoint(x, 1))(x_0) + grad_unrolled = jax.grad(lambda x: endpoint(x, 4))(x_0) + assert jnp.allclose(grad_rolled, grad_unrolled, rtol=1e-12, atol=1e-12) + with pytest.raises(ValueError, match="unroll"): + endpoint(x_0, 0) + + +def test_milstein_strong_order_gbm(): + keys = jax.random.split(jax.random.PRNGKey(2), 400) + levels = (16, 32, 64, 128) + errors = [] + for n in levels: + endpoints = jax.vmap( + lambda k, n=n: ( + solve_sde( + drift, + diffusion, + Milstein(), + 0.0, + T, + jnp.asarray(X_0), + key=k, + n_steps=n, + p=(MU, SIGMA), + ).xs + ) + )(keys) + exact = jax.vmap(lambda k, n=n: exact_endpoint(k, n))(keys) + errors.append(float(jnp.mean(jnp.abs(endpoints - exact)))) + slope = np.polyfit(np.log([T / n for n in levels]), np.log(errors), 1)[0] + assert 0.75 < slope < 1.25, (slope, errors) + + +def coarsen_sra1_noise(d_w, d_z, dt): + # Pasting adjacent steps: dW_c = dW_1 + dW_2 and the time-Wiener integral + # I_c = I_1 + I_2 + dt dW_1, then dZ back from I = (dt/2)(dW + dZ/sqrt(3)). + i_10 = 0.5 * dt * (d_w + d_z / np.sqrt(3.0)) + d_w_c = d_w[..., 0::2] + d_w[..., 1::2] + i_c = i_10[..., 0::2] + i_10[..., 1::2] + dt * d_w[..., 0::2] + d_z_c = np.sqrt(3.0) * (i_c / dt - d_w_c) + return d_w_c, d_z_c + + +def test_sra1_strong_order_additive_nonlinear(): + # No closed form conditions on (d_w, d_z) alone, so measure + # self-convergence against the finest grid on consistently coarsened + # noise; additive-noise EM on the same paths is strong order 1.0, so a + # slope well above 1 and the error ratio both separate the schemes. The + # observed slope may exceed the guaranteed 1.5 when the h^2 error terms + # dominate on the tested grids (exactly 2 for linear drift). + def cubic_drift(x): + return -x - x**3 + + n_fine = 256 + keys = jax.random.split(jax.random.PRNGKey(4), 200) + x_0 = jnp.asarray(X_0) + d_w, d_z = jax.vmap( + lambda k: SRA1().sample_noise( + x_0, k, n_fine, jnp.asarray(T / n_fine), x_0.dtype + ) + )(keys) + + def sra1_endpoints(d_w, d_z, n): + return jax.vmap( + lambda w, z: ( + solve_sde( + cubic_drift, + ou_diffusion, + SRA1(), + 0.0, + T, + x_0, + noise=(w, z), + n_steps=n, + ).xs + ) + )(d_w, d_z) + + reference = sra1_endpoints(d_w, d_z, n_fine) + level_w, level_z, n, dt = d_w, d_z, n_fine, T / n_fine + errors = {} + em_error = None + while n > 16: + level_w, level_z = coarsen_sra1_noise(level_w, level_z, dt) + n, dt = n // 2, 2.0 * dt + errors[n] = float( + jnp.mean(jnp.abs(sra1_endpoints(level_w, level_z, n) - reference)) + ) + if n == 16: + em = jax.vmap( + lambda w: ( + solve_sde( + cubic_drift, + ou_diffusion, + EulerMaruyama(), + 0.0, + T, + x_0, + noise=w, + n_steps=16, + ).xs + ) + )(level_w) + em_error = float(jnp.mean(jnp.abs(em - reference))) + ns = sorted(errors) + slope = np.polyfit(np.log([T / n for n in ns]), np.log([errors[n] for n in ns]), 1)[ + 0 + ] + assert 1.3 < slope < 2.5, (slope, errors) + assert errors[16] < 0.2 * em_error, (errors[16], em_error) + + +def test_zero_diffusion_sra1_is_ralston_rk2(): + # With diffusion = 0 the tableau collapses to the deterministic two-stage + # Ralston scheme, exact per step on dx = mu x: growth 1 + mu h + (mu h)^2/2. + n = 16 + h = T / n + sol = solve_sde( + lambda x: MU * x, + lambda x: 0.0 * x, + SRA1(), + 0.0, + T, + jnp.asarray(X_0), + key=jax.random.PRNGKey(0), + n_steps=n, + ) + growth = 1.0 + MU * h + 0.5 * (MU * h) ** 2 + assert jnp.abs(sol.xs - X_0 * growth**n) < 1e-12 + + +def test_zero_diffusion_milstein_matches_euler_ode(): + n = 16 + sde = solve_sde( + lambda x: MU * x, + lambda x: 0.0 * x, + Milstein(), + 0.0, + T, + jnp.asarray(X_0), + key=jax.random.PRNGKey(0), + n_steps=n, + save_at=SaveAt(steps=True), + ) + ode = solve_ode( + lambda x: MU * x, + Euler(), + 0.0, + T, + jnp.asarray(X_0), + dt_0=T / n, + max_steps=n, + save_at=SaveAt(steps=True), + ) + assert jnp.max(jnp.abs(sde.xs - ode.xs)) < 1e-14 diff --git a/tests/test_solvers_fixed.py b/tests/test_solvers_fixed.py index 27e0885..45805e8 100644 --- a/tests/test_solvers_fixed.py +++ b/tests/test_solvers_fixed.py @@ -2,18 +2,23 @@ import jax.numpy as jnp import jax.scipy.linalg as jsp_linalg import numpy as np +import pytest from kernels_reference import rk4_grid -from tinydiffeq import RK4, Euler, SaveAt, solve_ode +from tinydiffeq import RK4, Euler, IController, SaveAt, Tsit5, solve_ode def logistic_exact(x_0, t): return x_0 * jnp.exp(t) / (1.0 + x_0 * (jnp.exp(t) - 1.0)) -def test_linear_system_vs_expm(): - A = jnp.asarray([[0.0, 1.0], [-1.0, -0.3]]) - x_0 = jnp.asarray([1.0, 0.5]) +@pytest.mark.parametrize( + ("dtype", "euler_tol", "rk4_tol"), + [(jnp.float32, 4e-3, 1e-4), (jnp.float64, 2e-3, 1e-12)], +) +def test_linear_system_vs_expm(dtype, euler_tol, rk4_tol): + A = jnp.asarray([[0.0, 1.0], [-1.0, -0.3]], dtype) + x_0 = jnp.asarray([1.0, 0.5], dtype) T = 2.0 exact = jsp_linalg.expm(A * T) @ x_0 @@ -24,20 +29,26 @@ def f(x): euler = solve_ode(f, Euler(), 0.0, T, x_0, dt_0=T / n, max_steps=n) rk4 = solve_ode(f, RK4(), 0.0, T, x_0, dt_0=T / n, max_steps=n) assert bool(euler.ok) and bool(rk4.ok) + assert euler.xs.dtype == dtype and rk4.xs.dtype == dtype assert int(euler.num_steps) == int(euler.num_accepted) == n assert int(rk4.num_steps) == int(rk4.num_accepted) == n - assert jnp.max(jnp.abs(euler.xs - exact)) < 2e-3 - assert jnp.max(jnp.abs(rk4.xs - exact)) < 1e-12 + assert jnp.max(jnp.abs(euler.xs - exact)) < euler_tol + assert jnp.max(jnp.abs(rk4.xs - exact)) < rk4_tol -def test_logistic_closed_form(): - x_0 = jnp.asarray(0.1) +@pytest.mark.parametrize( + ("dtype", "tol"), + [(jnp.float32, 2e-5), (jnp.float64, 1e-9)], +) +def test_logistic_closed_form(dtype, tol): + x_0 = jnp.asarray(0.1, dtype) T = 3.0 n = 300 sol = solve_ode( lambda x: x * (1.0 - x), RK4(), 0.0, T, x_0, dt_0=T / n, max_steps=n ) - assert jnp.abs(sol.xs - logistic_exact(x_0, T)) < 1e-9 + assert sol.xs.dtype == dtype + assert jnp.abs(sol.xs - logistic_exact(x_0, jnp.asarray(T, dtype))) < tol def test_convergence_slopes(): @@ -59,6 +70,64 @@ def f(x): assert abs(slope - expected) < 0.3, (type(solver).__name__, slope) +def test_unroll_matches_rolled_and_requires_fixed_stepping(): + x_0 = jnp.asarray(0.1) + n = 16 + + def solve(unroll): + return solve_ode( + lambda x: x * (1.0 - x), + RK4(), + 0.0, + 1.0, + x_0, + dt_0=1.0 / n, + max_steps=n, + save_at=SaveAt(steps=True), + unroll=unroll, + ) + + rolled, unrolled = solve(1), solve(4) + assert jnp.array_equal(rolled.xs, unrolled.xs) + assert jnp.array_equal(rolled.ts, unrolled.ts) + grad_rolled = jax.grad( + lambda x: ( + solve_ode( + lambda v: v * (1.0 - v), RK4(), 0.0, 1.0, x, dt_0=1.0 / n, max_steps=n + ).xs + ) + )(x_0) + grad_unrolled = jax.grad( + lambda x: ( + solve_ode( + lambda v: v * (1.0 - v), + RK4(), + 0.0, + 1.0, + x, + dt_0=1.0 / n, + max_steps=n, + unroll=4, + ).xs + ) + )(x_0) + assert jnp.allclose(grad_rolled, grad_unrolled, rtol=1e-12, atol=1e-12) + with pytest.raises(ValueError, match="fixed stepping"): + solve_ode( + lambda x: -x, + Tsit5(), + 0.0, + 1.0, + x_0, + dt_0=0.1, + controller=IController(), + max_steps=32, + unroll=4, + ) + with pytest.raises(ValueError, match="at least 1"): + solve(0) + + def test_non_dividing_dt_0_lands_on_t_1(): x_0 = jnp.asarray(0.1) sol = solve_ode( diff --git a/tests/test_vmap_adaptive.py b/tests/test_vmap_adaptive.py new file mode 100644 index 0000000..e5f4e59 --- /dev/null +++ b/tests/test_vmap_adaptive.py @@ -0,0 +1,176 @@ +import jax +import jax.numpy as jnp +import pytest + +from tinydiffeq import ( + IController, + Rodas5P, + SaveAt, + Tsit5, + solve_ode, + solve_semi_explicit_dae, +) + +# Under vmap, a batched-predicate lax.cond lowers to a select that executes +# both branches, so without the scalar unvmap_all gates the adaptive loops +# would run every max_steps attempt slot for every lane. These tests pin the +# gates' semantics (per-lane adaptivity and results unchanged) and the work +# skipping itself. + +X_0S = jnp.linspace(0.5, 2.0, 8) + + +def adaptive_solve(x_0, max_steps=128, save_at=None): + return solve_ode( + lambda x, t, args, p: p * x * (1.0 - x), + Tsit5(), + 0.0, + 4.0, + x_0, + p=jnp.asarray(1.0), + dt_0=0.2, + controller=IController(), + max_steps=max_steps, + save_at=save_at, + ) + + +def test_vmapped_adaptive_matches_scalar_lanes(): + batched = jax.jit(jax.vmap(lambda x: adaptive_solve(x).xs))(X_0S) + stacked = jnp.stack([adaptive_solve(X_0S[i]).xs for i in range(len(X_0S))]) + assert jnp.allclose(batched, stacked, rtol=1e-12, atol=1e-12) + counts = jax.jit(jax.vmap(lambda x: adaptive_solve(x).num_steps))(X_0S) + scalar_counts = jnp.stack( + [adaptive_solve(X_0S[i]).num_steps for i in range(len(X_0S))] + ) + # Per-lane adaptivity: heterogeneous attempt counts survive batching. + assert jnp.array_equal(counts, scalar_counts) + assert int(counts.min()) != int(counts.max()) + + +def test_vmapped_adaptive_is_budget_invariant(): + small = jax.jit(jax.vmap(lambda x: adaptive_solve(x, max_steps=64).xs))(X_0S) + large = jax.jit(jax.vmap(lambda x: adaptive_solve(x, max_steps=512).xs))(X_0S) + assert jnp.array_equal(small, large) + + +def test_vmapped_adaptive_skips_frozen_tail(): + calls = [] + + def field(x, t, args, p): + jax.debug.callback(lambda: calls.append(1)) + return p * x * (1.0 - x) + + def solve(x_0): + return solve_ode( + field, + Tsit5(), + 0.0, + 4.0, + x_0, + p=jnp.asarray(1.0), + dt_0=0.2, + controller=IController(), + max_steps=256, + ).xs + + result = jax.block_until_ready(jax.vmap(solve)(X_0S)) + assert jnp.all(jnp.isfinite(result)) + # ~20 attempts x 7 stages of actual work; without the scalar gates every + # one of the 256 slots would evaluate the field (>1700 calls). + assert len(calls) < 600, len(calls) + + +def test_vmapped_adaptive_keeps_scalar_conds_in_jaxpr(): + jaxpr = str(jax.make_jaxpr(jax.vmap(lambda x: adaptive_solve(x).xs))(X_0S)) + assert "cond[" in jaxpr + assert "unvmap_all" in jaxpr + + +def test_grad_through_vmapped_adaptive_matches_scalar(): + def batched_loss(x_0s): + return jnp.sum(jax.vmap(lambda x: adaptive_solve(x).xs)(x_0s) ** 2) + + def scalar_loss(x_0s): + return sum(adaptive_solve(x_0s[i]).xs ** 2 for i in range(len(X_0S))) + + grad_batched = jax.jit(jax.grad(batched_loss))(X_0S) + grad_scalar = jax.grad(scalar_loss)(X_0S) + assert jnp.allclose(grad_batched, grad_scalar, rtol=1e-10, atol=1e-12) + tangent = jax.jvp(batched_loss, (X_0S,), (jnp.ones_like(X_0S),))[1] + assert jnp.isfinite(tangent) + + +def test_masked_residuals_on_vmapped_adaptive_output(): + # The collocation pattern from the docs: evaluate the pointwise residual + # on every padded row, zero the tail with `accepted`, normalize by the + # inert accepted count. fill="last" keeps padded rows finite, so the + # single where is safe for values and gradients. + def rollout(p): + return jax.vmap( + lambda x: solve_ode( + lambda x, t, args, p: p * x * (1.0 - x), + Tsit5(), + 0.0, + 4.0, + x, + p=p, + dt_0=0.2, + controller=IController(), + max_steps=64, + save_at=SaveAt(steps=True), + ) + )(X_0S) + + def residual(p): + sol = rollout(p) + rows = p * sol.xs * (1.0 - sol.xs) - (sol.xs - 1.0) + masked = jnp.where(sol.accepted, rows, 0.0) + count = jax.lax.stop_gradient(sol.accepted.sum()) + scaled = masked / jnp.sqrt(count.astype(masked.dtype)) + return jnp.where(sol.ok[:, None], scaled, jnp.inf).reshape(-1) + + p = jnp.asarray(1.2) + value = jax.jit(residual)(p) + assert value.shape == (len(X_0S) * 65,) + assert bool(jnp.all(jnp.isfinite(value))) + + # Padded rows contribute exactly zero: the flattened sum of squares + # equals the accepted-rows-only sum built lane by lane outside jit. + sol = rollout(p) + per_lane = 0.0 + for i in range(len(X_0S)): + lane_rows = p * sol.xs[i] * (1.0 - sol.xs[i]) - (sol.xs[i] - 1.0) + per_lane += float(jnp.sum(lane_rows[sol.accepted[i]] ** 2)) + count = float(sol.accepted.sum()) + assert jnp.allclose(jnp.sum(value**2), per_lane / count, rtol=1e-12) + + grad = jax.jit(jax.grad(lambda p: jnp.sum(residual(p) ** 2)))(p) + assert bool(jnp.isfinite(grad)) + + +@pytest.mark.parametrize("solver", [Tsit5(), Rodas5P()]) +def test_vmapped_adaptive_dae_matches_scalar_lanes(solver): + y_0s = jnp.linspace(0.8, 1.6, 4) + + def solve(y_0): + return solve_semi_explicit_dae( + lambda y, z, t, args, p: p * z, + lambda y, z: z - y, + solver, + 0.0, + 1.0, + y_0, + jnp.asarray(0.5), + p=jnp.asarray(-0.4), + dt_0=0.1, + controller=IController(), + max_steps=64, + save_at=SaveAt(t_1=True), + ) + + batched = jax.jit(jax.vmap(lambda y: solve(y).ys))(y_0s) + stacked = jnp.stack([solve(y_0s[i]).ys for i in range(len(y_0s))]) + assert jnp.allclose(batched, stacked, rtol=1e-12, atol=1e-12) + ok = jax.jit(jax.vmap(lambda y: solve(y).ok))(y_0s) + assert bool(jnp.all(ok)) diff --git a/uv.lock b/uv.lock index df761a5..7373b95 100644 --- a/uv.lock +++ b/uv.lock @@ -1227,7 +1227,7 @@ wheels = [ [[package]] name = "tinydiffeq" -version = "2.4.0" +version = "2.5.0" source = { editable = "." } dependencies = [ { name = "jax" },