tinydiffeq 0.1.0__tar.gz

tinydiffeq-0.1.0/.githooks/pre-commit

@@ -0,0 +1,5 @@
+#!/bin/sh
+set -e
+cd "$(git rev-parse --show-toplevel)"
+uv run ruff check .
+uv run ruff format --check .

tinydiffeq-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,29 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+      - run: uv sync --locked
+      - run: uv run ruff check .
+      - run: uv run ruff format --check .
+
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+          python-version: "3.13"
+      - run: uv sync --locked
+      - run: uv run pytest -q

tinydiffeq-0.1.0/.github/workflows/docs.yml

@@ -0,0 +1,39 @@
+name: Docs
+
+on:
+  push:
+    branches: [main]
+  workflow_dispatch:
+
+permissions:
+  contents: read
+  pages: write
+  id-token: write
+
+concurrency:
+  group: pages
+  cancel-in-progress: false
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+      - run: uv sync --group docs
+      - run: uv run mkdocs build --strict
+      - uses: actions/upload-pages-artifact@v3
+        with:
+          path: site
+
+  deploy:
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - id: deployment
+        uses: actions/deploy-pages@v4

tinydiffeq-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,36 @@
+name: Publish
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+      - run: uv build
+      - run: uvx twine check dist/*
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    environment:
+      name: pypi
+      url: https://pypi.org/p/tinydiffeq
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

tinydiffeq-0.1.0/.gitignore

@@ -0,0 +1,23 @@
+# Python
+__pycache__/
+*.py[cod]
+*.egg-info/
+.eggs/
+
+# Build artifacts
+build/
+dist/
+
+# Virtual environments
+.venv/
+
+# Docs build
+site/
+
+# Tooling caches
+.pytest_cache/
+.ruff_cache/
+.mypy_cache/
+
+# OS
+.DS_Store

tinydiffeq-0.1.0/.python-version

@@ -0,0 +1 @@
+3.13

tinydiffeq-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 HighDimensionalEconLab
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

tinydiffeq-0.1.0/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: tinydiffeq
+Version: 0.1.0
+Summary: Tiny differentiable ODE/SDE solvers for JAX — bounded-scan adaptive stepping, static shapes, jvp/vjp-safe
+Project-URL: Homepage, https://github.com/HighDimensionalEconLab/tinydiffeq
+Project-URL: Repository, https://github.com/HighDimensionalEconLab/tinydiffeq
+Project-URL: Documentation, https://highdimensionaleconlab.github.io/tinydiffeq/
+Project-URL: Issues, https://github.com/HighDimensionalEconLab/tinydiffeq/issues
+Author-email: Jesse Perla <jesseperla@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: autodiff,differential-equations,jax,ode,runge-kutta,sde
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Requires-Python: >=3.11
+Requires-Dist: jax>=0.7.0
+Description-Content-Type: text/markdown
+
+# tinydiffeq
+
+[![CI](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml)
+[![Docs](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/tinydiffeq/)
+[![PyPI](https://img.shields.io/pypi/v/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
+[![Python versions](https://img.shields.io/pypi/pyversions/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
+[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
+[![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 solvers for JAX: fixed-step Euler/RK4, adaptive
+Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
+adaptive stepping, so shapes are static, nothing recompiles as tolerances or
+curvature change, and every solve is differentiable in **both** forward and
+reverse mode — including reverse-over-forward, the pattern a
+Levenberg–Marquardt optimizer with geodesic acceleration needs when it
+differentiates through a rollout.
+
+This is a deliberately small, jvp/vjp-friendly subset of
+[diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
+states, stiff/implicit solvers, PID control, events, dense output, or
+checkpointed/backsolve adjoints. tinydiffeq's single runtime dependency is
+`jax`.
+
+## Install
+
+```bash
+uv add tinydiffeq
+```
+
+For GPU use, install the JAX accelerator build that matches your hardware,
+for example:
+
+```bash
+uv add tinydiffeq "jax[cuda13]"
+```
+
+## Minimal example
+
+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).
+
+```python
+import jax
+import jax.numpy as jnp
+from tinydiffeq import solve_ode, Tsit5, IController, SaveAt
+
+jax.config.update("jax_enable_x64", True)  # your call — the library never sets it
+
+
+def f(x, t, args, p):
+    return -p * x
+
+
+sol = solve_ode(
+    f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
+    p=jnp.asarray(1.3),
+    dt0=0.1,
+    controller=IController(rtol=1e-8, atol=1e-10),
+    max_steps=512,
+    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
+)                                                  # however many steps adapt
+print(sol.xs)   # states on the grid
+print(sol.ok)   # reached t1 within the max_steps budget?
+```
+
+## Gradients through the solve
+
+```python
+def endpoint(p):
+    return solve_ode(
+        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
+        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        max_steps=512,
+    ).xs
+
+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); the states differentiate fully through the RK stages. See the
+[docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
+contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
+semantics, and migration recipes from hand-rolled RK4/Tsit5 loops.
+
+## License
+
+MIT

tinydiffeq-0.1.0/README.md

@@ -0,0 +1,94 @@
+# tinydiffeq
+
+[![CI](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/ci.yml)
+[![Docs](https://github.com/HighDimensionalEconLab/tinydiffeq/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/tinydiffeq/)
+[![PyPI](https://img.shields.io/pypi/v/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
+[![Python versions](https://img.shields.io/pypi/pyversions/tinydiffeq.svg)](https://pypi.org/project/tinydiffeq/)
+[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/tinydiffeq)](https://github.com/HighDimensionalEconLab/tinydiffeq/blob/main/LICENSE)
+[![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 solvers for JAX: fixed-step Euler/RK4, adaptive
+Tsit5 with an integral step-size controller, and Euler–Maruyama for Itô SDEs.
+One bounded `lax.scan` of exactly `max_steps` iterations serves fixed and
+adaptive stepping, so shapes are static, nothing recompiles as tolerances or
+curvature change, and every solve is differentiable in **both** forward and
+reverse mode — including reverse-over-forward, the pattern a
+Levenberg–Marquardt optimizer with geodesic acceleration needs when it
+differentiates through a rollout.
+
+This is a deliberately small, jvp/vjp-friendly subset of
+[diffrax](https://docs.kidger.site/diffrax/). Use diffrax if you need pytree
+states, stiff/implicit solvers, PID control, events, dense output, or
+checkpointed/backsolve adjoints. tinydiffeq's single runtime dependency is
+`jax`.
+
+## Install
+
+```bash
+uv add tinydiffeq
+```
+
+For GPU use, install the JAX accelerator build that matches your hardware,
+for example:
+
+```bash
+uv add tinydiffeq "jax[cuda13]"
+```
+
+## Minimal example
+
+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).
+
+```python
+import jax
+import jax.numpy as jnp
+from tinydiffeq import solve_ode, Tsit5, IController, SaveAt
+
+jax.config.update("jax_enable_x64", True)  # your call — the library never sets it
+
+
+def f(x, t, args, p):
+    return -p * x
+
+
+sol = solve_ode(
+    f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
+    p=jnp.asarray(1.3),
+    dt0=0.1,
+    controller=IController(rtol=1e-8, atol=1e-10),
+    max_steps=512,
+    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),  # fixed output shape,
+)                                                  # however many steps adapt
+print(sol.xs)   # states on the grid
+print(sol.ok)   # reached t1 within the max_steps budget?
+```
+
+## Gradients through the solve
+
+```python
+def endpoint(p):
+    return solve_ode(
+        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
+        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        max_steps=512,
+    ).xs
+
+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); the states differentiate fully through the RK stages. See the
+[docs](https://highdimensionaleconlab.github.io/tinydiffeq/) for the design
+contracts: static shapes and `SaveAt`, AD through adaptive stepping, SDE key
+semantics, and migration recipes from hand-rolled RK4/Tsit5 loops.
+
+## License
+
+MIT

tinydiffeq-0.1.0/docs/adaptive_ad.md

@@ -0,0 +1,88 @@
+# Adaptive Stepping and AD
+
+tinydiffeq's central claim is that one adaptive solve is safely
+differentiable in forward mode, reverse mode, and reverse-over-forward. That
+holds because of three deliberate choices about **what is not
+differentiated**.
+
+## The controller is stop-gradiented
+
+The step-size controller (`IController`) computes its scaled error norm and
+next-step factor inside `stop_gradient`:
+
+> accept/reject is a non-differentiable branch either way, the gradient of
+> `E**(-1/order)` blows up at the exact-zero error of a flat-start policy,
+> and the `d(dt)/dtheta` term only slides sample points along the visited
+> trajectory — irrelevant to a residual that must vanish at every state.
+
+The states themselves remain fully differentiable through the RK stages: the
+gradient you get is the derivative of the numerical flow map *for the step
+pattern actually taken*. This is the same convention diffrax uses.
+
+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 `max(E, 1e-12)` floor) the backward pass is NaN
+from iteration one. `tests/test_ad.py::test_grad_finite_on_flat_field` pins
+this.
+
+## The horizon clip is the growth guard
+
+Every attempt is clipped so it cannot step past `t1`, 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 `factormax = 5`, a flat field would otherwise reach `dt ≈ t1/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 `factormax × remaining horizon`.
+
+One refinement over a bare `min(dt, remaining)`: when the remaining horizon
+is within `max_steps × eps` of the desired step, the step is stretched to
+land on `t1` exactly. Summing `n` rounded steps of `(t1 - t0)/n` can leave
+`t` one accumulated ulp short of `t1`, and without the stretch that sliver
+would cost an extra iteration a `max_steps = n` budget doesn't have.
+
+## Interpolation knots are non-differentiable
+
+`SaveAt(ts=...)` brackets each query with `searchsorted` — integer indices,
+no gradient. This is consistent with the stop-gradiented controller: the
+term excluded is again "the knots slide along the trajectory as parameters
+change". Values differentiate fully through the bracketing states and
+derivatives (`xL`, `xR`, `fL`, `fR`).
+
+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
+
+- `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`.
+- `jax.vmap` over `x0` or `p` gives genuinely per-lane adaptivity: each lane
+  accepts/rejects independently through the masked scan, and batched results
+  equal the individual solves exactly.
+
+## What to watch
+
+- Reductions over `SaveAt(steps=True)` rows are **discontinuous** in the
+  inputs: the number of duplicate rows changes when an accept flips to a
+  reject. Consume steps mode with residuals that vanish at every state (so
+  duplicates are harmless), or reduce to `xs[-1]`.
+- 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.

tinydiffeq-0.1.0/docs/api.md

@@ -0,0 +1,35 @@
+# API Reference
+
+## Solve functions
+
+::: tinydiffeq.solve_ode
+
+::: tinydiffeq.solve_sde
+
+## Solvers
+
+::: tinydiffeq.Euler
+
+::: tinydiffeq.RK4
+
+::: tinydiffeq.Tsit5
+
+::: tinydiffeq.EulerMaruyama
+
+## Step-size controllers
+
+::: tinydiffeq.ConstantStepSize
+
+::: tinydiffeq.IController
+
+## Output selection and results
+
+::: tinydiffeq.SaveAt
+
+::: tinydiffeq.Solution
+
+## Utilities
+
+::: tinydiffeq.hermite_interpolate
+
+::: tinydiffeq.cumulative_trapezoid

tinydiffeq-0.1.0/docs/index.md

@@ -0,0 +1,125 @@
+# tinydiffeq
+
+`tinydiffeq` is a deliberately tiny set of differentiable ODE/SDE integrators
+for JAX: fixed-step Euler and RK4, adaptive Tsit5 with an integral step-size
+controller, and fixed-step Euler–Maruyama for Itô SDEs. Everything runs
+inside one bounded `lax.scan` with static shapes, and every solve is
+differentiable in **both** forward and reverse mode — including
+reverse-over-forward, the pattern a Levenberg–Marquardt optimizer with
+geodesic acceleration needs when it differentiates through a rollout.
+
+It is a jvp/vjp-friendly subset of
+[diffrax](https://docs.kidger.site/diffrax/). **Use diffrax instead if you
+need any of:**
+
+- pytree states (tinydiffeq states are single arrays, scalar or vector)
+- stiff or implicit solvers
+- PID step-size control
+- events, root-finding, or backward-time integration
+- dense output / continuous interpolation objects
+- checkpointed or backsolve adjoints for long horizons
+
+tinydiffeq ships only the O(max_steps)-memory bounded-scan approach, because
+that is the one that composes cleanly with `jax.jvp`, `jax.vjp`, `jax.vmap`,
+and reverse-over-forward without custom adjoint machinery.
+
+## Install
+
+```bash
+uv add tinydiffeq
+```
+
+For accelerator use, install the JAX build matching your hardware alongside
+it, for example:
+
+```bash
+uv add tinydiffeq "jax[cuda13]"
+```
+
+## Vector-field interface
+
+The vector field may take one to four positional arguments — always in this
+order:
+
+```python
+f(x)                # autonomous, closes over everything
+f(x, t)
+f(x, t, args)
+f(x, t, args, p)
+```
+
+- `x` is the array state (scalar or vector).
+- `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 `x0` 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.
+
+## Minimal example
+
+```python
+import jax
+import jax.numpy as jnp
+from tinydiffeq import solve_ode, Tsit5, IController, SaveAt
+
+jax.config.update("jax_enable_x64", True)  # your call, not the library's
+
+
+def f(x, t, args, p):
+    return -p * x
+
+
+sol = solve_ode(
+    f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0),
+    p=jnp.asarray(1.3),
+    dt0=0.1,
+    controller=IController(rtol=1e-8, atol=1e-10),
+    max_steps=512,
+    saveat=SaveAt(ts=jnp.linspace(0.0, 2.0, 21)),
+)
+sol.xs  # (21,) states on the grid, however many internal steps were taken
+sol.ok  # False if max_steps ran out before t1
+```
+
+Gradients go straight through the solve:
+
+```python
+def endpoint(p):
+    return solve_ode(
+        f, Tsit5(), 0.0, 2.0, jnp.asarray(1.0), p=p,
+        dt0=0.1, controller=IController(rtol=1e-10, atol=1e-12),
+        max_steps=512,
+    ).xs
+
+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
+
+- **`dt0` is required.** There is no initial-step heuristic.
+- **Forward time only**: `t1 > t0`.
+- **Never poisons.** `sol.ok` reports whether `t1` was reached; callers that
+  want diverging residuals do `jnp.where(sol.ok, sol.xs, jnp.inf)`.
+- **`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
+  `jnp.result_type(x0, float)`; float32 problems stay float32 even under
+  x64.
+- Solvers, controllers, `SaveAt`, and `Solution` are frozen dataclasses
+  registered as pytrees: numeric fields (tolerances, grids, `dt0`, `x0`) are
+  data leaves, so changing them never recompiles.
+
+Read next: [Static Shapes](static_shapes.md) for the bounded-scan design and
+`SaveAt`, [Adaptive Stepping and AD](adaptive_ad.md) for what is and is not
+differentiated, [SDEs](sde.md), and
+[Migration](migration.md) if you are replacing hand-rolled RK4/Tsit5 loops.

tinydiffeq-0.1.0/docs/llms.txt

@@ -0,0 +1,12 @@
+# tinydiffeq
+
+> Tiny differentiable ODE/SDE solvers for JAX. Fixed-step (Euler, RK4) and adaptive (Tsit5 + integral controller) explicit Runge-Kutta inside one bounded lax.scan of exactly max_steps iterations: static shapes, one compilation across tolerances/curvature/initial conditions, and solves differentiable in both forward and reverse mode (including reverse-over-forward) with O(max_steps) memory. Array states only (scalar or vector). SaveAt selects endpoint, cubic-Hermite interpolation onto a fixed grid, or raw padded step rows. solve_sde is fixed-step Euler-Maruyama with presampled diagonal noise keyed on a PRNG key. A deliberately small subset of diffrax: use diffrax for pytree states, stiff/implicit solvers, PID control, events, dense output, and adjoint methods.
+
+## Docs
+
+- [Home — positioning, vector-field signature convention f(x, t, args, p), minimal examples](https://highdimensionaleconlab.github.io/tinydiffeq/)
+- [Static shapes — the bounded-scan design, SaveAt as the shape contract, duplicate rows vs inf padding, why nothing recompiles](https://highdimensionaleconlab.github.io/tinydiffeq/static_shapes/)
+- [Adaptive stepping and AD — stop-gradiented controller rationale, horizon-clip growth guard, non-differentiable interpolation knots, double-where NaN safety](https://highdimensionaleconlab.github.io/tinydiffeq/adaptive_ad/)
+- [SDEs — Euler-Maruyama orders, fixed-noise key semantics, shared-path strong-convergence testing, why SaveAt(ts) raises](https://highdimensionaleconlab.github.io/tinydiffeq/sde/)
+- [Migration — exact recipes replacing hand-rolled rk4_grid/tsit5_free loops, poisoning via jnp.where(sol.ok, ...), documented behavior changes](https://highdimensionaleconlab.github.io/tinydiffeq/migration/)
+- [API reference — solve_ode, solve_sde, solvers, controllers, SaveAt, Solution, hermite_interpolate, cumulative_trapezoid](https://highdimensionaleconlab.github.io/tinydiffeq/api/)