openscvx 0.5.3.dev16__tar.gz → 0.5.3.dev17__tar.gz

{openscvx-0.5.3.dev16/openscvx.egg-info → openscvx-0.5.3.dev17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev16
+Version: 0.5.3.dev17
 Summary: A general Python-based successive convexification implementation which uses a JAX backend.
 Author-email: Chris Hayner and Griffin Norris <haynec@uw.edu>
 License: Apache Software License

openscvx-0.5.3.dev17/docs/UnderTheHood/batching_jit_grad.md

@@ -0,0 +1,134 @@
+# Batching, JIT, and Grad with `Problem.solve_jax()`
+
+`Problem` exposes two solve entry points. `solve()` is the familiar
+Python-loop driver — real-time prints, wall-clock `time_limit`, `continuous`
+mode, populated per-iteration history. `solve_jax()` is its JAX-pure
+sibling: it drives the same fused `iteration_fn` body inside a
+`lax.while_loop` and returns an `OptimizationResults`
+(`openscvx/algorithms/optimization_results.py`) pytree, so it composes with
+`jax.vmap`, `jax.jit`, and `jax.grad`.
+
+```python
+problem = ox.Problem(..., solver={"backend": "qpax"})
+problem.initialize()
+
+# Familiar interactive solve — unchanged: prints, history, time limit.
+result = problem.solve()
+
+# JAX-pure single solve — silent, no per-iteration history, batchable.
+result = problem.solve_jax()
+
+# Batched — standard jax.vmap, no library-specific API.
+batched = jax.vmap(problem.solve_jax, in_axes=(0, 0, None))(
+    x_initial_stack, x_final_stack, params
+)
+
+# JIT-compile once, solve forever (e.g. MPC inner loop).
+fast_solve = jax.jit(problem.solve_jax)
+```
+
+A worked example lives at `examples/abstract/brachistochrone_batched.py`
+— four brachistochrone problems with shifted starting x-coordinates,
+solved in parallel via `jax.vmap(problem.solve_jax)`.
+
+## When to choose `solve()` vs. `solve_jax()`
+
+| Use case | Entry point |
+|---|---|
+| Interactive use, real-time prints, plotting per iteration | `solve()` |
+| Wall-clock `time_limit` or `continuous=True` | `solve()` |
+| Per-iteration trajectories / weights / diagnostics | `solve()` |
+| Batched solve over many initial conditions or parameters | `solve_jax()` |
+| JIT-compile once, solve forever (MPC inner loop, scenario sweeps) | `solve_jax()` |
+| `jax.grad` through the solver (best-effort, untested) | `solve_jax()` |
+
+The split exists to make user intent explicit — a single dispatched
+`solve()` would have to infer from arguments which path the caller wants,
+and silent routing failures (like `continuous=True` quietly skipping the
+Python loop) are exactly the failure class the two-method design exists to
+prevent. See `plans/jax-pure-solve.md`'s Decision Log for the longer
+discussion.
+
+## What `solve_jax()` returns
+
+`solve_jax()` returns the same `OptimizationResults` type as `solve()`,
+but built via
+`OptimizationResults.from_final_state(state, problem=...)` instead of
+`from_history(history, final_state, ...)`. The differences:
+
+* **Per-iteration history is empty.** `X = [state.x]` and `U = [state.u]`
+  are single-element lists so `result.x` / `result.u` continue to return
+  the final iterate; every `*_history` field is `[]`. List growth doesn't
+  fit inside `lax.while_loop`. If you need per-iteration trajectories, use
+  `solve()` or wrap `solve_jax()` in `lax.scan` manually.
+* **Post-process fields stay `None`.** `t_full`, `x_full`, `u_full`,
+  `cost`, `ctcs_violation` are populated only by `post_process()`. They're
+  outside the JAX pytree's children surface, so a batched `solve_jax`
+  result doesn't force every consumer to handle `None` leaves. If you want
+  to post-process, call `post_process()` per batch element after the
+  batched solve.
+* **`converged` is a `jnp.bool_`** under the JAX-pure path (a `(B,)` array
+  under vmap), not a Python `bool` like `solve()` returns. Most uses
+  (`if result.converged: ...`) work either way.
+
+## `solve_jax()` arguments
+
+```python
+problem.solve_jax(
+    x_initial=None,      # boundary-condition pin (full unified state vector,
+                         #     ``jnp.nan`` at non-Fix entries — see
+                         #     ``AlgorithmState.from_settings``); falls back
+                         #     to the default from settings
+    x_final=None,        # terminal pin, same conventions
+    parameters=None,     # parameters dict for this solve; falls back to
+                         #     ``self._parameters``
+    *,
+    max_iters=None,      # SCP iteration cap; non-default rebuilds the
+                         #     cached ``lax.while_loop`` closure (one extra
+                         #     trace; subsequent calls at the same cap hit
+                         #     the cache)
+)
+```
+
+Positional kwargs (rather than a single `inputs` pytree) keep
+`jax.vmap(problem.solve_jax, in_axes=(0, 0, None))` ergonomic;
+multi-argument gradient uses `jax.grad(..., argnums=(0, 1))`.
+
+## Caveats
+
+### CVXPy under `vmap` is sequential
+
+The CVXPy backend's `iteration_callback`
+(`openscvx/solvers/cvxpy_ptr_solver.py`) host-calls CVXPy through
+`jax.pure_callback` with `vmap_method="sequential"`. Host
+CVXPy is not thread-safe, so a `jax.vmap(problem.solve_jax)` over the
+CVXPy backend runs `B` sequential CVXPy solves. The QPAX and Moreau
+backends are pure JAX end-to-end and run in parallel under vmap.
+
+### Moreau warm-start is bypassed
+
+The Moreau backend's `_warm_start` carry is host-side mutable state that
+`lax.while_loop` doesn't thread. Both `solve()` and `solve_jax()` cold-start
+the inner Moreau solve every SCP iteration (see the Moreau module
+docstring and `plans/jax-pure-solve.md`'s Decision Log 2026-05-27).
+Restoring warm-start to the SCP loop requires threading `(x, z, s)`
+through an `AlgorithmState.moreau_carry` field — a Future Extension.
+
+### `jax.grad` is best-effort, untested
+
+The QPAX backend's `solve_qp` is not differentiable (its convergence flag
+costs reverse-mode autodiff). The CVXPy backend's `pure_callback` is
+non-differentiable by default. End-to-end gradient validation through
+`solve_jax` is a follow-up — `jax.grad(loss_fn)(x0)` will run without
+erroring under QPAX (`solve_qp_primal` would be `custom_vjp`, but the
+backend currently uses `solve_qp`), but the gradient's correctness is not
+yet pinned down by a test against finite differences.
+
+### Convergence under `vmap` with mixed-rate elements
+
+`lax.while_loop` continues while *any* batch element still needs
+iterations — converged elements would otherwise keep receiving body calls
+while their peers iterate. `make_solve_loop`'s body selects the unchanged
+state for converged elements (`jax.tree.map(jnp.where, ...)`), so a
+batched solve agrees with the single-problem `solve_jax` on each element.
+The cost is per-iteration work × slowest-element iteration count.

{openscvx-0.5.3.dev16 → openscvx-0.5.3.dev17}/docs/UnderTheHood/lowering_architecture.md

@@ -188,8 +188,13 @@ JAX export, used only by `post_process()` — so a compiled problem holds **one
 `iteration_fn` plus one propagation export**, not three discretization /
 propagation exports stitched together in Python. Because `iteration_fn` is a
 registered pytree in and out, it also composes with `jax.jit`; `make_solve_loop`
-wraps it in `lax.while_loop` as the primitive for a future fully-JAX-driven
-solve.
+wraps it in `lax.while_loop` to drive the **JAX-pure** `Problem.solve_jax()`
+entry point, which composes with `jax.vmap`, `jax.jit`, and `jax.grad` (see
+`docs/UnderTheHood/batching_jit_grad.md`). The familiar
+`Problem.solve()` keeps its Python-loop shape — real-time prints, wall-clock
+`time_limit`, populated per-iteration history — and is the right entry
+point for interactive use. `solve_jax()` is the right entry point for
+batched / JIT'd / differentiable workflows.
 
 Picking a backend at construction time:
 

openscvx-0.5.3.dev17/examples/abstract/brachistochrone_batched.py

@@ -0,0 +1,126 @@
+"""Batched brachistochrone via ``jax.vmap(problem.solve_jax)``.
+
+Solves four brachistochrone problems in parallel — each with a different
+starting x-coordinate — using the JAX-pure
+:meth:`~openscvx.problem.Problem.solve_jax` entry point. ``solve_jax``
+returns an :class:`~openscvx.algorithms.optimization_results.OptimizationResults`
+pytree, so ``jax.vmap`` over the boundary-condition pins is the standard JAX
+composition pattern; no library-specific batching API is needed.
+
+Run with ``python examples/abstract/brachistochrone_batched.py``.
+
+Notes:
+
+* Under the QPAX backend used here, vmap'd subproblem solves run in
+  parallel (no host callback). The CVXPy backend is single-threaded under
+  vmap because :func:`jax.pure_callback` uses ``vmap_method="sequential"``
+  for thread safety — see :meth:`solve_jax`'s docstring.
+* The Moreau warm-start is bypassed on the JAX-pure path (its carry is
+  host-side state that ``lax.while_loop`` doesn't thread). Use QPAX or
+  CVXPy when you want batched solves.
+"""
+
+import os
+import sys
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+current_dir = os.path.dirname(os.path.abspath(__file__))
+grandparent_dir = os.path.dirname(os.path.dirname(current_dir))
+sys.path.append(grandparent_dir)
+
+import openscvx as ox
+from openscvx import Problem
+
+
+def build_problem(n: int = 8):
+    g = 9.81
+
+    position = ox.State("position", shape=(2,))
+    position.max = np.array([10.0, 10.0])
+    position.min = np.array([0.0, 0.0])
+    position.initial = np.array([0.0, 10.0])
+    position.final = [10.0, 5.0]
+
+    velocity = ox.State("velocity", shape=(1,))
+    velocity.max = np.array([10.0])
+    velocity.min = np.array([0.0])
+    velocity.initial = np.array([0.0])
+    velocity.final = [("free", 10.0)]
+
+    theta = ox.Control("theta", shape=(1,))
+    theta.max = np.array([100.5 * jnp.pi / 180])
+    theta.min = np.array([0.0])
+    theta.guess = np.linspace(5 * jnp.pi / 180, 100.5 * jnp.pi / 180, n).reshape(-1, 1)
+
+    states = [position, velocity]
+    controls = [theta]
+
+    dynamics = {
+        "position": ox.Concat(
+            velocity[0] * ox.Sin(theta[0]),
+            -velocity[0] * ox.Cos(theta[0]),
+        ),
+        "velocity": g * ox.Cos(theta[0]),
+    }
+
+    constraint_exprs = []
+    for state in states:
+        constraint_exprs.extend([ox.ctcs(state <= state.max), ox.ctcs(state.min <= state)])
+
+    time = ox.Time(
+        initial=0.0,
+        final=("minimize", 2.0),
+        min=0.0,
+        max=2.0,
+        uniform_time_grid=True,
+    )
+
+    return Problem(
+        dynamics=dynamics,
+        states=states,
+        controls=controls,
+        time=time,
+        constraints=constraint_exprs,
+        N=n,
+        float_dtype="float64",
+        algorithm={
+            "autotuner": "ConstantProximalWeight",
+            "lam_prox": 1e0,
+            "lam_cost": 6e-1,
+            "ep_tr": 1e-5,
+            "ep_vb": 1e-5,
+            "ep_vc": 1e-9,
+            "k_max": 30,
+        },
+        solver={"backend": "qpax"},
+    )
+
+
+if __name__ == "__main__":
+    problem = build_problem()
+    problem.initialize()
+
+    # The default boundary-condition pin (the full unified state vector with
+    # ``jnp.nan`` at non-Fix entries — see
+    # :meth:`AlgorithmState.from_settings`).
+    default_pin = problem.state.x_init_pin
+
+    # Stack four ICs by varying the starting x-coordinate (component 0 of
+    # the unified state vector is the x-position).
+    shifts = jnp.array([0.0, 0.3, 0.6, 0.9])
+    x_initial_stack = jnp.stack([default_pin.at[0].set(default_pin[0] + s) for s in shifts])
+
+    # Bare ``jax.vmap`` — the library exposes no batching wrapper of its own.
+    batched_solve = jax.vmap(problem.solve_jax, in_axes=(0, None, None))
+    results = batched_solve(x_initial_stack, None, None)
+
+    # ``results`` is a batched ``OptimizationResults`` pytree — every child
+    # (``X``, ``U``, ``t_final``, ``converged``, ...) has a leading batch axis.
+    print(f"Batched solve over {x_initial_stack.shape[0]} initial conditions:")
+    print(f"  result.x.shape:   {results.x.shape}")
+    print(f"  result.u.shape:   {results.u.shape}")
+    print(f"  result.t_final:   {np.asarray(results.t_final).reshape(-1)}")
+    print(f"  result.converged: {np.asarray(results.converged)}")

{openscvx-0.5.3.dev16 → openscvx-0.5.3.dev17}/mkdocs.yml

@@ -197,4 +197,5 @@ nav:
       - 'Foundations/discretization.md'
     - Under the Hood:
       - 'UnderTheHood/vectorization_and_vmapping.md'
-      - 'UnderTheHood/lowering_architecture.md'
+      - 'UnderTheHood/lowering_architecture.md'
+      - 'UnderTheHood/batching_jit_grad.md'

{openscvx-0.5.3.dev16 → openscvx-0.5.3.dev17}/openscvx/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.3.dev16'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev16')
+__version__ = version = '0.5.3.dev17'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev17')
 
-__commit_id__ = commit_id = 'gb2479bf05'
+__commit_id__ = commit_id = 'gb2143958f'

{openscvx-0.5.3.dev16 → openscvx-0.5.3.dev17}/openscvx/algorithms/optimization_results.py

@@ -1,10 +1,17 @@
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import TYPE_CHECKING, Any, Optional, Union
 
+import jax
+import jax.numpy as jnp
 import numpy as np
 
+if TYPE_CHECKING:
+    from openscvx.algorithms.base import AlgorithmHistory, AlgorithmState
+    from openscvx.problem import Problem
 
+
+@jax.tree_util.register_pytree_node_class
 @dataclass
 class OptimizationResults:
     """
@@ -44,6 +51,30 @@ class OptimizationResults:
             Added by propagate_trajectory_results.
         plotting_data (dict[str, Any]): Flexible storage for plotting and application data.
 
+    !!! note "JAX pytree registration"
+        ``OptimizationResults`` is a registered JAX pytree, so the output of
+        :meth:`~openscvx.problem.Problem.solve_jax` composes with
+        ``jax.vmap`` / ``jax.jit`` / ``jax.grad``. The pytree *children* are
+        ``converged``, ``t_final``, ``nodes``, ``trajectory``, ``X``, ``U``,
+        and every ``*_history`` field — the data the user actually consumes.
+        Post-process fields (``t_full``, ``x_full``, ``u_full``, ``cost``,
+        ``ctcs_violation``), ``plotting_data``, and the internal ``_states`` /
+        ``_controls`` metadata are stashed in the treedef's *aux* instead, so
+        a batched ``solve_jax`` doesn't force callers to handle ``None``
+        leaves and ``post_process()`` outputs stay outside the vmap surface
+        (post-process per element after the batched solve).
+
+        The two construction paths produce different histories:
+
+        * :meth:`from_history` (``solve()``) populates the per-iteration
+          ``X`` / ``U`` / ``*_history`` lists from the Python loop's
+          :class:`~openscvx.algorithms.base.AlgorithmHistory`.
+        * :meth:`from_final_state` (``solve_jax()``) wraps only the final
+          iterate: ``X = [state.x]`` and ``U = [state.u]`` (single-element
+          lists so ``result.x`` / ``result.u`` continue to return the final
+          iterate), and every ``*_history`` field is empty — per-iteration
+          history requires Python-side list growth that can't run inside
+          ``lax.while_loop``.
 
     !!! note "For Developers"
         The ``metadata={"npz": ...}`` parameter on each field below is a built-in feature of
@@ -66,6 +97,30 @@ class OptimizationResults:
     _OPT_ARRAY = "optional_array"
     _OPT_SCALAR = "optional_scalar"
 
+    # Pytree children: fields whose contents flow through ``jax.vmap`` /
+    # ``jax.jit`` / ``jax.grad``. Excludes ``_states`` / ``_controls`` (Python
+    # symbolic metadata), ``plotting_data`` (user scratch), and the
+    # post-process fields (``None`` until ``post_process()`` runs and not part
+    # of the ``solve_jax`` vmap surface — see the class docstring).
+    _PYTREE_CHILDREN = (
+        "converged",
+        "t_final",
+        "nodes",
+        "trajectory",
+        "X",
+        "U",
+        "discretization_history",
+        "J_tr_history",
+        "J_vb_history",
+        "J_vc_history",
+        "TR_history",
+        "VC_history",
+        "lam_prox_history",
+        "actual_reduction_history",
+        "pred_reduction_history",
+        "acceptance_ratio_history",
+    )
+
     # Core optimization results
     converged: bool = field(metadata={"npz": "scalar"})
     t_final: float = field(metadata={"npz": "scalar"})
@@ -138,6 +193,183 @@ class OptimizationResults:
         """Initialize the results object."""
         pass
 
+    # ------------------------------------------------------------------
+    # JAX pytree registration
+    # ------------------------------------------------------------------
+
+    def tree_flatten(self):
+        """Split into JAX-traceable children and host-side aux.
+
+        See :attr:`_PYTREE_CHILDREN` for the children list. Aux carries the
+        symbolic metadata, scratch dict, and post-process fields verbatim so
+        the round-trip ``tree_unflatten(*tree_flatten(self))`` preserves the
+        instance.
+        """
+        children = tuple(getattr(self, name) for name in self._PYTREE_CHILDREN)
+        aux = (
+            self._states,
+            self._controls,
+            self.t_full,
+            self.x_full,
+            self.u_full,
+            self.cost,
+            self.ctcs_violation,
+            self.plotting_data,
+        )
+        return children, aux
+
+    @classmethod
+    def tree_unflatten(cls, aux, children):
+        kwargs = dict(zip(cls._PYTREE_CHILDREN, children))
+        (
+            _states,
+            _controls,
+            t_full,
+            x_full,
+            u_full,
+            cost,
+            ctcs_violation,
+            plotting_data,
+        ) = aux
+        return cls(
+            **kwargs,
+            _states=_states,
+            _controls=_controls,
+            t_full=t_full,
+            x_full=x_full,
+            u_full=u_full,
+            cost=cost,
+            ctcs_violation=ctcs_violation,
+            plotting_data=plotting_data,
+        )
+
+    # ------------------------------------------------------------------
+    # Constructors
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_final_state(
+        cls,
+        state: "AlgorithmState",
+        *,
+        problem: "Problem",
+    ) -> "OptimizationResults":
+        """JAX-pure construction from the final SCP iterate.
+
+        Used by :meth:`~openscvx.problem.Problem.solve_jax`. Every leaf is a
+        ``jnp.ndarray`` derived from ``state`` via static slicing, so the
+        result composes with ``jax.vmap`` / ``jax.jit`` / ``jax.grad`` over a
+        batched ``state``. ``X = [state.x]`` and ``U = [state.u]`` are
+        single-element lists so ``result.x`` / ``result.u`` continue to
+        return the final iterate; every ``*_history`` field is empty — the
+        Python-loop ``.solve()`` is the path that populates them.
+        """
+        settings = problem.settings
+        symbolic = problem.symbolic
+        algorithm = problem._algorithm
+
+        has_impulsive_controls = any(
+            control.parameterization == "impulsive" for control in symbolic.controls
+        )
+
+        nodes_dict: dict[str, jnp.ndarray] = {}
+        for sym_state in symbolic.states:
+            state_nodes = state.x[..., sym_state._slice]
+            if has_impulsive_controls:
+                bc = jnp.asarray(settings.sim.x.initial[sym_state._slice])
+                state_nodes = state_nodes.at[..., 0, :].set(bc)
+            nodes_dict[sym_state.name] = state_nodes
+
+        for control in symbolic.controls:
+            nodes_dict[control.name] = state.u[..., control._slice]
+
+        # Match the host-path ``t_final`` shape: a 1-vector along the time
+        # slice of the last node (vmap'd to ``(B, 1)``).
+        t_final = state.x[..., -1, :][..., settings.sim.time_slice]
+
+        converged = (
+            (state.J_tr < algorithm.ep_tr)
+            & (state.J_vb < algorithm.ep_vb)
+            & (state.J_vc < algorithm.ep_vc)
+        )
+
+        return cls(
+            converged=converged,
+            t_final=t_final,
+            nodes=nodes_dict,
+            trajectory={},
+            _states=symbolic.states_prop,
+            _controls=symbolic.controls,
+            X=[state.x],
+            U=[state.u],
+            discretization_history=[],
+            J_tr_history=[],
+            J_vb_history=[],
+            J_vc_history=[],
+            TR_history=[],
+            VC_history=[],
+            lam_prox_history=[],
+            actual_reduction_history=[],
+            pred_reduction_history=[],
+            acceptance_ratio_history=[],
+        )
+
+    @classmethod
+    def from_history(
+        cls,
+        history: "AlgorithmHistory",
+        final_state: "AlgorithmState",
+        *,
+        problem: "Problem",
+        converged: bool,
+    ) -> "OptimizationResults":
+        """Host-side construction from the Python-loop iteration history.
+
+        Used by :meth:`~openscvx.problem.Problem.solve` to package the full
+        per-iteration log produced by the Python ``while`` loop. Symmetric
+        with :meth:`from_final_state` — same return type, populated history.
+        """
+        settings = problem.settings
+        symbolic = problem.symbolic
+
+        has_impulsive_controls = any(
+            control.parameterization == "impulsive" for control in symbolic.controls
+        )
+
+        x_arr = np.asarray(final_state.x)
+        u_arr = np.asarray(final_state.u)
+
+        nodes_dict: dict[str, np.ndarray] = {}
+        for sym_state in symbolic.states:
+            state_nodes = x_arr[:, sym_state._slice].copy()
+            if has_impulsive_controls:
+                state_nodes[0] = settings.sim.x.initial[sym_state._slice]
+            nodes_dict[sym_state.name] = state_nodes
+
+        for control in symbolic.controls:
+            nodes_dict[control.name] = u_arr[:, control._slice]
+
+        return cls(
+            converged=converged,
+            t_final=x_arr[:, settings.sim.time_slice][-1],
+            nodes=nodes_dict,
+            trajectory={},
+            _states=symbolic.states_prop,
+            _controls=symbolic.controls,
+            X=history.X,
+            U=history.U,
+            discretization_history=history.V_history,
+            J_tr_history=float(final_state.J_tr),
+            J_vb_history=float(final_state.J_vb),
+            J_vc_history=float(final_state.J_vc),
+            TR_history=history.TR,
+            VC_history=history.VC,
+            lam_prox_history=history.lam_prox.copy(),
+            actual_reduction_history=history.actual_reduction.copy(),
+            pred_reduction_history=history.pred_reduction.copy(),
+            acceptance_ratio_history=history.acceptance_ratio.copy(),
+        )
+
     def update_plotting_data(self, **kwargs: Any) -> None:
         """
         Update the plotting data with additional information.

{openscvx-0.5.3.dev16 → openscvx-0.5.3.dev17}/openscvx/algorithms/scvx/iteration.py

@@ -361,8 +361,25 @@ def make_solve_loop(
             return (state.k <= k_max) & jnp.logical_not(_converged(state, ep_tr, ep_vb, ep_vc))
 
         def body(state: AlgorithmState) -> AlgorithmState:
+            # Under ``jax.vmap`` the ``lax.while_loop`` keeps running until
+            # every batch element has converged; without a freeze, the body
+            # would keep mutating already-converged elements (their iterates
+            # drift through repeated subproblem solves, and the autotuner
+            # would keep advancing ``lam_prox`` / ``lam_cost``). Selecting
+            # ``state`` for converged elements pins them to their first
+            # post-convergence iterate, so a batched solve agrees with the
+            # single-problem ``solve_jax`` on each element.
+            is_converged = _converged(state, ep_tr, ep_vb, ep_vc)
             next_state, _ = iteration_fn(state, params)
-            return next_state
+
+            def freeze(nxt, prev):
+                # ``is_converged`` is scalar (single-problem) or shape ``(B,)``
+                # (vmap'd); reshape with trailing 1-axes so it broadcasts over
+                # each leaf's remaining dims.
+                mask_shape = is_converged.shape + (1,) * (nxt.ndim - is_converged.ndim)
+                return jnp.where(is_converged.reshape(mask_shape), prev, nxt)
+
+            return jax.tree.map(freeze, next_state, state)
 
         return jax.lax.while_loop(cond, body, state)