openscvx 0.5.3.dev18__tar.gz → 0.5.3.dev20__tar.gz

{openscvx-0.5.3.dev18/openscvx.egg-info → openscvx-0.5.3.dev20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev18
+Version: 0.5.3.dev20
 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.dev18 → openscvx-0.5.3.dev20}/docs/UnderTheHood/batching_jit_grad.md

@@ -49,6 +49,81 @@ 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.
 
+## Batched, disk-cached solves with `solve_batched()`
+
+`jax.vmap(solve_jax)` and `solve_batched()` both run `B` SCP solves over
+stacked boundary conditions; they differ in **who owns the batch axis**, and
+that one difference is the whole point.
+
+* `jax.vmap(solve_jax)` — *the caller* owns the axis. Maximally idiomatic,
+  composes with `grad` / `scan`, but **in-process JIT only**: `jax.export` has
+  no outer `vmap` rule (`call_exported` cannot be batched), so an exported
+  artifact can never be re-vmapped, and every fresh process re-traces and
+  re-compiles the entire loop.
+* `solve_batched(x0_stack, xf_stack)` — *the library* owns the axis. The vmap
+  is applied **inside** the method, before any export, so the whole loop lowers
+  to ordinary batched XLA ops and serializes cleanly. Under
+  `save_compiled=True` it is written to the solver cache on the first run and
+  deserialized on later processes — skipping the XLA compile that
+  `jax.vmap(solve_jax)` pays on every launch.
+
+```python
+problem = ox.Problem(..., solver={"backend": "qpax"})
+problem.settings.sim.save_compiled = True
+problem.initialize()
+
+# First process: traces, exports, writes <cache>/compiled_solve_batched_<hash>.jax
+# Later processes: deserialize-and-.call, no compile.
+batched = problem.solve_batched(x0_stack, xf_stack, params)
+# batched.x.shape == (B, N, n_states) — same pytree jax.vmap(solve_jax) returns.
+```
+
+Pick `solve_batched` when **cross-process cold-start dominates** — short-lived
+processes (CI sweeps, notebook restarts, deployed workers) that otherwise pay
+the compile every launch. Pick `jax.vmap(solve_jax)` for everything in-program,
+where you own the transforms and want `grad` / `scan` composition. With
+`save_compiled=False`, `solve_batched` is just an in-process batched wrapper and
+the two produce identical results.
+
+The exported artifact is closed and deployable, not composable: it is the
+`call_exported`-has-no-outer-`vmap`-rule wall again. Use `solve_jax` for
+user-driven `jax.vmap` / `jax.grad`; `solve_batched` for a cached batched solve.
+
+### CVXPy cannot be exported
+
+Export requires every op in the loop — the backend solve included — to lower to
+serializable XLA. CVXPy's `iteration_callback` is a `jax.pure_callback`, and
+`jax.export` cannot serialize host callbacks. So `PTRSolver.exportable` is
+`False` for CVXPy (and `True` for QPAX / Moreau), and
+`solve_batched(save_compiled=True)` on the CVXPy backend **raises a teaching
+error** pointing at QPAX / Moreau rather than silently degrading to an uncached
+in-process solve. With `save_compiled=False`, CVXPy still runs `B` sequential
+in-process solves.
+
+### Cache invalidation is correctness, not just freshness
+
+The exported loop closes over far more than the symbolic problem does. Caching
+on the symbolic hash alone would silently deserialize a **stale artifact** when
+any of that extra state changes — a wrong-answer bug, not a cache miss. So
+`get_solve_batched_cache_path` (`openscvx/utils/caching.py`) extends the
+symbolic hash with everything additionally baked into the loop, each runtime
+object contributing through its own `_hash_into` (the same protocol the symbolic
+layer uses):
+
+* the convex backend class + its `solver_args`,
+* the algorithm + autotuner + convergence thresholds + initial weights,
+* the discretizer scheme (class + hold type + ODE solver),
+* the state/control scaling matrices `inv_S_x` / `inv_S_u`,
+* the fixed batch size `B` (the artifact is exported at one `B`), and
+* the JAX version (exported artifacts are not guaranteed to survive an
+  incompatible jax bump).
+
+Change any of these — backend, a tolerance, `k_max`, a state bound — and you get
+a different cache path, so a stale artifact is never reused.
+
+`examples/abstract/brachistochrone_batched.py` shows both paths side by side;
+re-running it is itself the cross-process demo.
+
 ## What `solve_jax()` returns
 
 `solve_jax()` returns the same `OptimizationResults` type as `solve()`,

{openscvx-0.5.3.dev18 → openscvx-0.5.3.dev20}/examples/abstract/brachistochrone_batched.py

@@ -1,20 +1,30 @@
-"""Batched brachistochrone via ``jax.vmap(problem.solve_jax)``.
+"""Batched brachistochrone — two ways to solve four problems at once.
 
 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``.
+starting x-coordinate — and contrasts the two batching entry points:
+
+* ``jax.vmap(problem.solve_jax)`` — the caller owns the batch axis. Maximally
+  idiomatic, composes with ``grad`` / ``scan``, but in-process JIT only: every
+  fresh process re-traces and re-compiles the whole SCP loop.
+* ``problem.solve_batched(x0_stack, xf_stack)`` — the library owns the batch
+  axis, so the whole vmapped loop is a single function that ``jax.export`` can
+  serialize. Under ``save_compiled=True`` it is written to the solver cache on
+  the first run and deserialized on later runs, skipping that compile. Reach
+  for it when cross-process cold-start dominates (CI sweeps, short-lived
+  workers); reach for ``jax.vmap(solve_jax)`` when you stay inside one program.
+
+Run with ``python examples/abstract/brachistochrone_batched.py``. Re-running is
+itself the cross-process demo: the second invocation deserializes the cached
+``solve_batched`` artifact instead of compiling it.
 
 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.
+  for thread safety — see :meth:`solve_jax`'s docstring. That same host
+  callback is why ``solve_batched(save_compiled=True)`` refuses CVXPy: it
+  cannot be exported.
 * 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.
@@ -113,14 +123,33 @@ if __name__ == "__main__":
     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.
+    # --- caller owns the batch axis: jax.vmap(solve_jax) ---------------------
+    # Bare ``jax.vmap`` — composes like any JAX transform, in-process only.
     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"jax.vmap(solve_jax) 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)}")
+
+    # --- library owns the batch axis: solve_batched, disk-cached -------------
+    # The same batched solve, but ``solve_batched`` owns the vmap so the whole
+    # loop is one exportable artifact. Under ``save_compiled=True`` it lands in
+    # the solver cache on the first run; re-run this script (a fresh process) to
+    # see it deserialized instead of recompiled. ``xf`` is shared across the
+    # batch, so broadcast the default terminal pin to the same leading axis.
+    export_problem = build_problem()
+    export_problem.settings.sim.save_compiled = True
+    export_problem.initialize()
+
+    x_term_stack = jnp.broadcast_to(export_problem.state.x_term_pin, x_initial_stack.shape)
+    batched = export_problem.solve_batched(x_initial_stack, x_term_stack)
+
+    print(f"\nsolve_batched (save_compiled=True) over {x_initial_stack.shape[0]} ICs:")
+    print(f"  result.x.shape:   {batched.x.shape}")
+    print(f"  result.t_final:   {np.asarray(batched.t_final).reshape(-1)}")
+    print(f"  result.converged: {np.asarray(batched.converged)}")