nlls-gram 0.1.0__tar.gz

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

@@ -0,0 +1,33 @@
+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
+    strategy:
+      fail-fast: false
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+          python-version: ${{ matrix.python-version }}
+      - run: uv sync --locked
+      - run: uv run pytest -q

nlls_gram-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

nlls_gram-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/nlls-gram
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+      - uses: pypa/gh-action-pypi-publish@release/v1

nlls_gram-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

nlls_gram-0.1.0/.python-version

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

nlls_gram-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.

nlls_gram-0.1.0/PKG-INFO

@@ -0,0 +1,104 @@
+Metadata-Version: 2.4
+Name: nlls-gram
+Version: 0.1.0
+Summary: Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax NNX models
+Project-URL: Homepage, https://github.com/HighDimensionalEconLab/nlls_gram
+Project-URL: Repository, https://github.com/HighDimensionalEconLab/nlls_gram
+Project-URL: Documentation, https://highdimensionaleconlab.github.io/nlls_gram/
+Project-URL: Issues, https://github.com/HighDimensionalEconLab/nlls_gram/issues
+Author-email: Jesse Perla <jesseperla@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Keywords: flax,jax,levenberg-marquardt,nnx,nonlinear-least-squares,optimization
+Classifier: Development Status :: 3 - Alpha
+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: flax>=0.10.7
+Requires-Dist: jax>=0.7.0
+Requires-Dist: optax>=0.2.4
+Description-Content-Type: text/markdown
+
+# nlls_gram
+
+[![CI](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/ci.yml)
+[![Docs](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/nlls_gram/)
+[![PyPI](https://img.shields.io/pypi/v/nlls-gram.svg)](https://pypi.org/project/nlls-gram/)
+[![Python versions](https://img.shields.io/pypi/pyversions/nlls-gram.svg)](https://pypi.org/project/nlls-gram/)
+[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/nlls_gram)](https://github.com/HighDimensionalEconLab/nlls_gram/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)
+
+Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax
+NNX models.
+
+`GramLevenbergMarquardt` minimizes `||r(theta)||^2` for a residual defined over an
+`nnx.Module`, following the optax/nnx `init`/`update` protocol so that steps apply
+through `nnx.Optimizer(model, optax.identity(), wrt=...)`. For overparameterized
+systems (many more parameters `p` than residual rows `n`) it factors the small
+`n x n` gram (dual) system instead of the `p x p` normal equations.
+
+## Install
+
+```bash
+pip install nlls-gram
+```
+
+## Minimal example
+
+Fit `y = a * exp(b * x)` to noise-free data generated from `(a, b) = (2, -1)`:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+from flax import nnx
+
+from nlls_gram import GramLevenbergMarquardt
+
+jax.config.update("jax_enable_x64", True)
+
+
+class ExpModel(nnx.Module):
+    def __init__(self, a, b):
+        self.a = nnx.Param(jnp.asarray(a))
+        self.b = nnx.Param(jnp.asarray(b))
+
+    def __call__(self, x):
+        return self.a * jnp.exp(self.b * x)
+
+
+x = jnp.linspace(0.0, 2.0, 20)
+y = 2.0 * jnp.exp(-1.0 * x)
+
+model = ExpModel(a=1.0, b=0.0)
+solver = GramLevenbergMarquardt(lambda m, batch: m(batch[0]) - batch[1])
+optimizer = nnx.Optimizer(model, optax.identity(), wrt=nnx.Param)
+lm_state = solver.init()
+
+
+@jax.jit
+def train_step(graphdef, state, lm_state, batch):
+    m, opt = nnx.merge(graphdef, state)
+    updates, lm_state, info = solver.update(m, lm_state, batch)
+    opt.update(m, updates)
+    return lm_state, info, nnx.state((m, opt))
+
+
+graphdef, state = nnx.split((model, optimizer))
+for _ in range(50):
+    lm_state, info, state = train_step(graphdef, state, lm_state, (x, y))
+nnx.update((model, optimizer), state)
+
+print(model.a[...], model.b[...])  # ~2.0, ~-1.0
+```
+
+## Documentation
+
+Full docs: https://highdimensionaleconlab.github.io/nlls_gram/
+
+## License
+
+MIT

nlls_gram-0.1.0/README.md

@@ -0,0 +1,80 @@
+# nlls_gram
+
+[![CI](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/ci.yml/badge.svg)](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/ci.yml)
+[![Docs](https://github.com/HighDimensionalEconLab/nlls_gram/actions/workflows/docs.yml/badge.svg)](https://highdimensionaleconlab.github.io/nlls_gram/)
+[![PyPI](https://img.shields.io/pypi/v/nlls-gram.svg)](https://pypi.org/project/nlls-gram/)
+[![Python versions](https://img.shields.io/pypi/pyversions/nlls-gram.svg)](https://pypi.org/project/nlls-gram/)
+[![License: MIT](https://img.shields.io/github/license/HighDimensionalEconLab/nlls_gram)](https://github.com/HighDimensionalEconLab/nlls_gram/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)
+
+Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax
+NNX models.
+
+`GramLevenbergMarquardt` minimizes `||r(theta)||^2` for a residual defined over an
+`nnx.Module`, following the optax/nnx `init`/`update` protocol so that steps apply
+through `nnx.Optimizer(model, optax.identity(), wrt=...)`. For overparameterized
+systems (many more parameters `p` than residual rows `n`) it factors the small
+`n x n` gram (dual) system instead of the `p x p` normal equations.
+
+## Install
+
+```bash
+pip install nlls-gram
+```
+
+## Minimal example
+
+Fit `y = a * exp(b * x)` to noise-free data generated from `(a, b) = (2, -1)`:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+from flax import nnx
+
+from nlls_gram import GramLevenbergMarquardt
+
+jax.config.update("jax_enable_x64", True)
+
+
+class ExpModel(nnx.Module):
+    def __init__(self, a, b):
+        self.a = nnx.Param(jnp.asarray(a))
+        self.b = nnx.Param(jnp.asarray(b))
+
+    def __call__(self, x):
+        return self.a * jnp.exp(self.b * x)
+
+
+x = jnp.linspace(0.0, 2.0, 20)
+y = 2.0 * jnp.exp(-1.0 * x)
+
+model = ExpModel(a=1.0, b=0.0)
+solver = GramLevenbergMarquardt(lambda m, batch: m(batch[0]) - batch[1])
+optimizer = nnx.Optimizer(model, optax.identity(), wrt=nnx.Param)
+lm_state = solver.init()
+
+
+@jax.jit
+def train_step(graphdef, state, lm_state, batch):
+    m, opt = nnx.merge(graphdef, state)
+    updates, lm_state, info = solver.update(m, lm_state, batch)
+    opt.update(m, updates)
+    return lm_state, info, nnx.state((m, opt))
+
+
+graphdef, state = nnx.split((model, optimizer))
+for _ in range(50):
+    lm_state, info, state = train_step(graphdef, state, lm_state, (x, y))
+nnx.update((model, optimizer), state)
+
+print(model.a[...], model.b[...])  # ~2.0, ~-1.0
+```
+
+## Documentation
+
+Full docs: https://highdimensionaleconlab.github.io/nlls_gram/
+
+## License
+
+MIT

nlls_gram-0.1.0/docs/index.md

@@ -0,0 +1,78 @@
+# nlls_gram
+
+Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax
+NNX models.
+
+`GramLevenbergMarquardt` minimizes `||r(theta)||^2` for a residual defined over an
+`nnx.Module`, following the optax/nnx `init`/`update` protocol so that steps apply
+through `nnx.Optimizer(model, optax.identity(), wrt=...)`. For overparameterized
+systems (many more parameters `p` than residual rows `n`) it factors the small
+`n x n` gram (dual) system instead of the `p x p` normal equations.
+
+## Install
+
+```bash
+pip install nlls-gram
+```
+
+## Minimal example
+
+Fit `y = a * exp(b * x)` to noise-free data generated from `(a, b) = (2, -1)`:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+from flax import nnx
+
+from nlls_gram import GramLevenbergMarquardt
+
+jax.config.update("jax_enable_x64", True)
+
+
+class ExpModel(nnx.Module):
+    def __init__(self, a, b):
+        self.a = nnx.Param(jnp.asarray(a))
+        self.b = nnx.Param(jnp.asarray(b))
+
+    def __call__(self, x):
+        return self.a * jnp.exp(self.b * x)
+
+
+# residual_fn(model, batch) -> residual array; the solver minimizes its sum of squares
+def residual_fn(model, batch):
+    x, y = batch
+    return model(x) - y
+
+
+x = jnp.linspace(0.0, 2.0, 20)
+y = 2.0 * jnp.exp(-1.0 * x)
+
+model = ExpModel(a=1.0, b=0.0)
+solver = GramLevenbergMarquardt(residual_fn, init_damping=1e-2)
+optimizer = nnx.Optimizer(model, optax.identity(), wrt=nnx.Param)
+lm_state = solver.init()
+
+
+# The solver does not jit internally; wrap the train step yourself.
+@jax.jit
+def train_step(graphdef, state, lm_state, batch):
+    m, opt = nnx.merge(graphdef, state)
+    updates, lm_state, info = solver.update(m, lm_state, batch)
+    opt.update(m, updates)
+    return lm_state, info, nnx.state((m, opt))
+
+
+graphdef, state = nnx.split((model, optimizer))
+for _ in range(50):
+    lm_state, info, state = train_step(graphdef, state, lm_state, (x, y))
+nnx.update((model, optimizer), state)
+
+print(model.a[...], model.b[...])  # ~2.0, ~-1.0
+```
+
+## API reference
+
+::: nlls_gram.GramLevenbergMarquardt
+
+::: nlls_gram.flat_residual

nlls_gram-0.1.0/mkdocs.yml

@@ -0,0 +1,40 @@
+site_name: nlls_gram
+site_description: Gram/dual-form Levenberg-Marquardt nonlinear least-squares for JAX/Flax NNX models
+site_url: https://highdimensionaleconlab.github.io/nlls_gram/
+repo_url: https://github.com/HighDimensionalEconLab/nlls_gram
+repo_name: HighDimensionalEconLab/nlls_gram
+
+theme:
+  name: material
+  features:
+    - content.code.copy
+    - navigation.top
+  palette:
+    - scheme: default
+      primary: indigo
+      toggle:
+        icon: material/brightness-7
+        name: Switch to dark mode
+    - scheme: slate
+      primary: indigo
+      toggle:
+        icon: material/brightness-4
+        name: Switch to light mode
+
+plugins:
+  - search
+  - mkdocstrings:
+      handlers:
+        python:
+          options:
+            show_root_heading: true
+            show_source: true
+
+markdown_extensions:
+  - pymdownx.highlight:
+      anchor_linenums: true
+  - pymdownx.superfences
+  - admonition
+
+nav:
+  - Home: index.md

nlls_gram-0.1.0/pyproject.toml

@@ -0,0 +1,63 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "nlls-gram"
+version = "0.1.0"
+description = "Gram/dual-form Levenberg-Marquardt nonlinear least-squares solvers for JAX/Flax NNX models"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.11"
+authors = [{ name = "Jesse Perla", email = "jesseperla@gmail.com" }]
+keywords = [
+    "levenberg-marquardt",
+    "nonlinear-least-squares",
+    "jax",
+    "flax",
+    "nnx",
+    "optimization",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Mathematics",
+]
+dependencies = [
+    "flax>=0.10.7",
+    "jax>=0.7.0",
+    "optax>=0.2.4",
+]
+
+[project.urls]
+Homepage = "https://github.com/HighDimensionalEconLab/nlls_gram"
+Repository = "https://github.com/HighDimensionalEconLab/nlls_gram"
+Documentation = "https://highdimensionaleconlab.github.io/nlls_gram/"
+Issues = "https://github.com/HighDimensionalEconLab/nlls_gram/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/nlls_gram"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py311"
+
+[tool.ruff.lint]
+select = ["E", "F", "I", "UP", "B"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+
+[dependency-groups]
+dev = [
+    "pytest>=8",
+    "ruff>=0.15.17",
+]
+docs = [
+    "mkdocs-material>=9.5",
+    "mkdocstrings[python]>=0.26",
+]

nlls_gram-0.1.0/src/nlls_gram/__init__.py

@@ -0,0 +1,16 @@
+"""Gram/dual-form Levenberg-Marquardt nonlinear least-squares for JAX/Flax NNX models.
+
+GramLevenbergMarquardt minimizes ||r(theta)||^2 for an nnx.Module-defined residual,
+following the optax/nnx init/update protocol so steps apply through
+nnx.Optimizer(model, optax.identity(), wrt=...). For overparameterized systems
+(p parameters >> n residual rows) it factors the small n x n gram (dual) system.
+"""
+
+from nlls_gram.gram_lm import (
+    GramLevenbergMarquardt,
+    LMInfo,
+    LMState,
+    flat_residual,
+)
+
+__all__ = ["GramLevenbergMarquardt", "LMState", "LMInfo", "flat_residual"]

nlls_gram-0.1.0/src/nlls_gram/gram_lm.py

@@ -0,0 +1,94 @@
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+from flax import nnx
+from jax.flatten_util import ravel_pytree
+
+# optax/nnx-style least-squares and root-finding solvers: each class exposes
+# init() -> state and update(model, state, batch) -> (updates, state, info), where
+# updates are parameter *increments* structured like nnx.split(model, wrt, ...)[1],
+# meant to be applied through nnx.Optimizer(model, optax.identity(), wrt=...).
+# Classes do not jit internally -- wrap the caller's train step in jax.jit. All
+# hyperparameters are static Python scalars; all data-dependent control flow is
+# traced (jnp.where / lax.while_loop), so a rejected step returns zero updates
+# rather than branching.
+
+
+# Functionalize a model into a flat parameter vector and a residual function of it:
+# residual_fn(model, batch) becomes residual_flat(theta) with theta the raveled
+# wrt-filtered state (non-wrt state captured and passed through unchanged).
+def flat_residual(model, wrt, residual_fn, batch):
+    graphdef, diff_state, rest = nnx.split(model, wrt, ...)
+    theta, unravel = ravel_pytree(diff_state)
+
+    def residual_flat(th):
+        m = nnx.merge(graphdef, unravel(th), rest)
+        return jnp.ravel(residual_fn(m, batch))
+
+    return theta, unravel, residual_flat
+
+
+class LMState(NamedTuple):
+    damping: jax.Array
+
+
+class LMInfo(NamedTuple):
+    loss: jax.Array  # min(old, new) sum of squared residuals
+    accepted: jax.Array
+    damping: jax.Array  # post-update damping
+
+
+# Classic Marquardt damping on min ||R(theta)||^2: accept the step iff the sum of
+# squared residuals decreases, multiplying the damping by damping_decrease on
+# acceptance and damping_increase on rejection. solve_method picks which linear
+# system is factored for the identical step:
+#   "gram":   step = -J' (J J' + damping I_n)^{-1} R   (n x n dual; right for n << p)
+#   "normal": step = -(J'J + damping I_p)^{-1} J' R    (p x p; right for p <~ n)
+class GramLevenbergMarquardt:
+    def __init__(
+        self,
+        residual_fn,
+        *,
+        init_damping=1e-3,
+        damping_decrease=0.5,
+        damping_increase=4.0,
+        solve_method="gram",
+        wrt=nnx.Param,
+    ):
+        if solve_method not in ("gram", "normal"):
+            raise ValueError(f"unknown solve_method: {solve_method}")
+        self.residual_fn = residual_fn
+        self.init_damping = init_damping
+        self.damping_decrease = damping_decrease
+        self.damping_increase = damping_increase
+        self.solve_method = solve_method
+        self.wrt = wrt
+
+    def init(self):
+        return LMState(jnp.asarray(self.init_damping))
+
+    def update(self, model, state, batch):
+        theta, unravel, residual_flat = flat_residual(
+            model, self.wrt, self.residual_fn, batch
+        )
+        resid = residual_flat(theta)
+        J = jax.jacrev(residual_flat)(theta)
+        if self.solve_method == "gram":
+            step = -J.T @ jnp.linalg.solve(
+                J @ J.T + state.damping * jnp.eye(J.shape[0]), resid
+            )
+        else:
+            step = -jnp.linalg.solve(
+                J.T @ J + state.damping * jnp.eye(J.shape[1]), J.T @ resid
+            )
+        resid_new = residual_flat(theta + step)
+        improved = jnp.sum(resid_new**2) < jnp.sum(resid**2)
+        updates_flat = jnp.where(improved, step, jnp.zeros_like(step))
+        damping = jnp.where(
+            improved,
+            state.damping * self.damping_decrease,
+            state.damping * self.damping_increase,
+        )
+        loss = jnp.minimum(jnp.sum(resid_new**2), jnp.sum(resid**2))
+        return unravel(updates_flat), LMState(damping), LMInfo(loss, improved, damping)

nlls_gram-0.1.0/tests/test_gram_lm.py

@@ -0,0 +1,50 @@
+import jax
+import jax.numpy as jnp
+import optax
+from flax import nnx
+
+from nlls_gram import GramLevenbergMarquardt
+
+jax.config.update("jax_enable_x64", True)
+
+
+class ExpModel(nnx.Module):
+    def __init__(self, a, b):
+        self.a = nnx.Param(jnp.asarray(a))
+        self.b = nnx.Param(jnp.asarray(b))
+
+    def __call__(self, x):
+        return self.a * jnp.exp(self.b * x)
+
+
+def residual_fn(model, batch):
+    x, y = batch
+    return model(x) - y
+
+
+def test_recovers_known_parameters():
+    a_true, b_true = 2.0, -1.0
+    x = jnp.linspace(0.0, 2.0, 20)
+    y = a_true * jnp.exp(b_true * x)
+
+    model = ExpModel(a=1.0, b=0.0)
+    solver = GramLevenbergMarquardt(residual_fn, init_damping=1e-2)
+    optimizer = nnx.Optimizer(model, optax.identity(), wrt=nnx.Param)
+    lm_state = solver.init()
+
+    @jax.jit
+    def train_step(graphdef, state, lm_state, batch):
+        m, opt = nnx.merge(graphdef, state)
+        updates, lm_state, info = solver.update(m, lm_state, batch)
+        opt.update(m, updates)
+        return lm_state, info, nnx.state((m, opt))
+
+    graphdef, state = nnx.split((model, optimizer))
+    info = None
+    for _ in range(50):
+        lm_state, info, state = train_step(graphdef, state, lm_state, (x, y))
+    nnx.update((model, optimizer), state)
+
+    assert float(info.loss) < 1e-10
+    assert jnp.allclose(model.a[...], a_true, atol=1e-4)
+    assert jnp.allclose(model.b[...], b_true, atol=1e-4)