openscvx 0.5.3.dev23__tar.gz → 0.5.3.dev24__tar.gz

{openscvx-0.5.3.dev23/openscvx.egg-info → openscvx-0.5.3.dev24}/PKG-INFO

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

@@ -60,7 +60,7 @@ that one difference is the whole point.
   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
+* `solve_batched(x_initial=x0_batch)` — *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
@@ -74,10 +74,119 @@ 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 = problem.solve_batched(x_initial=x0_batch, parameters=params)
 # batched.x.shape == (B, N, n_states) — same pytree jax.vmap(solve_jax) returns.
 ```
 
+### One rule: add a leading batch axis to whatever varies
+
+`solve_batched` takes exactly `solve_jax`'s arguments. Every batchable input
+has a *declared* (unbatched) shape the library knows — `(n_x,)` for the
+boundary pins, `(N, n_x)` / `(N, n_u)` for the trajectory guesses, the
+parameter's stored shape for each parameter, and `()` for the SCP
+hyperparameters. An argument that matches its declared shape is **shared** by
+every batch element; one with exactly one extra leading axis is **batched**
+along it. Rank against the declared shape decides — never the leading-axis
+value — so a shared `(B, n)` matrix parameter is never misread as batched.
+(`B` in these shape patterns is plain notation for the batch size, not
+syntax: it is never passed anywhere — like `jax.vmap`, the library reads it
+off the leading axis of whatever you batched and cross-checks that all
+batched inputs agree.)
+
+```python
+results = problem.solve_batched(
+    x_initial=x0_batch,                  # (B, n_x): rank = declared+1 -> batched
+    parameters={
+        "gate_center": centers,          # (B, 3) vs declared (3,)  -> batched
+        "gate_radius": 0.5,              # declared shape           -> shared
+    },
+)
+```
+
+For the rare entry the rank rule cannot express — e.g. batching a declared
+rank-1 parameter element-wise — pass an explicit `in_axes` in `jax.vmap`'s
+vocabulary, including its prefix semantics (`0` = batched along the leading
+axis, `None` = shared; a prefix applies to everything below it; partial specs
+fall back to the rank rule):
+
+```python
+results = problem.solve_batched(
+    parameters={"weights": w},           # (B,) is also w's declared shape...
+    in_axes={"parameters": {"weights": 0}},  # ...so say "batched" explicitly
+)
+
+results = problem.solve_batched(parameters=params, in_axes={"parameters": 0})
+# prefix: every passed parameter is batched, exactly like jax.vmap's
+# pytree-prefix in_axes; a bare in_axes=0 batches every passed input.
+```
+
+One deliberate divergence from `jax.vmap`: the *default* is "infer by rank",
+not `0`-everything — an all-shared batched solve has no batch axis to find,
+and inference is what makes mixed shared/batched dicts ergonomic.
+
+### Hyperparameter sweeps: the `algorithm` dict
+
+Solve-time mirrors construction-time: the `Problem` constructor configures
+the SCP algorithm with a dict (`algorithm={"ep_tr": 1e-5, "lam_prox": 1.0}`),
+and the solve methods take the same box with the same names as per-solve,
+batchable inputs. Every `AlgorithmState` field is a runtime input riding the
+state pytree, and the `algorithm` dict reaches any of them by name —
+convergence thresholds (`ep_tr` / `ep_vb` / `ep_vc`), penalty weights
+(`lam_prox` / `lam_vc` / `lam_cost`), the seed iterate, autotuner
+hyperparameters. There is no whitelist: the field declaration is the
+registration, and the `AlgorithmState` attribute docs are the reference.
+(`max_iters` survives as a named kwarg — sugar for `algorithm={"k_max":
+...}`; passing both raises, as does an entry that shadows any other named
+kwarg's field. Construction-only keys of the constructor dict — `autotuner`,
+`t_max` — raise pointing back at the constructor.)
+
+* On `solve_jax`, each value matches the field's shape, or is a scalar
+  broadcast to it: `solve_jax(algorithm={"ep_tr": 1e-6, "lam_prox": 2.0})`
+  retraces nothing.
+* On `solve_batched`, each value follows the rank rule against the field's
+  shape, plus two **fill** forms for array-shaped fields: a scalar (shared)
+  or a `(B,)` vector (one scalar per element, broadcast to the field shape).
+  Exact shapes win when both parse; `in_axes={"algorithm": {name: 0}}`
+  forces the per-element reading. All against one cached artifact:
+
+```python
+B = 8  # batch size. Just an int for building the arrays below — there is no
+       # B argument; solve_batched infers it from the leading-axis lengths.
+sweep = problem.solve_batched(algorithm={"ep_tr": jnp.logspace(-6, -3, B)})
+weights = problem.solve_batched(algorithm={"lam_prox": jnp.array([0.5, 1.0, 4.0])})
+budgets = problem.solve_batched(max_iters=jnp.array([5, 10, 20]))  # per-element caps
+```
+
+The batch runs until its slowest element converges or exhausts its own cap;
+finished elements are frozen, not re-iterated, so each element matches the
+corresponding single `solve_jax` run.
+
+**New autotuners batch for free.** An autotuner's tunable knobs follow the
+same derivation: declare them as fields on a `HyperParams` subclass — bare
+annotations, no decorator; subclassing applies the frozen-dataclass transform
+and the JAX pytree registration, and the annotation picks the dtype (`int`
+gets the iteration counter's, `float` the problem's) — assign an instance to
+`self.hyper`, and read `state.hyper.<field>` inside `update_weights` (see the
+`AutotuningBase` docstring). The declaration is the registration — each field
+immediately works as a `solve_jax` override, a `solve_batched` sweep target,
+and a runtime input of the exported artifact, with zero edits anywhere else:
+
+```python
+from openscvx.algorithms import AutotuningBase, HyperParams
+
+class MyHyper(HyperParams):
+    ramp: float = 2.0
+
+class MyAutotuner(AutotuningBase):
+    def __init__(self, ramp: float = 2.0):
+        self.hyper = MyHyper(ramp=ramp)
+
+    def update_weights(self, state, ...):
+        ...state.hyper.ramp...
+
+results = problem.solve_batched(algorithm={"ramp": jnp.linspace(1.5, 3.0, B)})
+```
+
 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,
@@ -111,15 +220,21 @@ 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 algorithm + autotuner + 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 fixed batch size `B` (the artifact is exported at one `B`),
+* the per-parameter shared/batched split (baked into the `jax.vmap` program),
+  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.
+Change any of these — backend, a penalty weight, a state bound, which
+parameters are batched — and you get a different cache path, so a stale
+artifact is never reused. Everything the `algorithm` dict can reach — convergence
+thresholds, `max_iters`, weights, autotuner hyperparameters — is deliberately
+*not* in the key: those are runtime inputs on the state pytree, so one
+artifact serves every setting of them.
 
 `examples/abstract/brachistochrone_batched.py` shows both paths side by side;
 re-running it is itself the cross-process demo.
@@ -158,10 +273,15 @@ problem.solve_jax(
     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)
+    x_guess=None,        # (N, n_x) state trajectory warm-start
+    u_guess=None,        # (N, n_u) control trajectory warm-start
+    max_iters=None,      # SCP iteration cap; sugar for algorithm={"k_max": ...};
+                         #     falls back to the configured k_max
+    algorithm=None,      # any AlgorithmState field by name (ep_tr, lam_prox,
+                         #     ...) — the constructor's algorithm-dict names;
+                         #     values match the field shape or fill it from a
+                         #     scalar. Runtime inputs on the state pytree — no
+                         #     value forces a retrace
 )
 ```
 

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/examples/abstract/brachistochrone_batched.py

@@ -6,7 +6,7 @@ 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
+* ``problem.solve_batched(x_initial=...)`` — 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
@@ -140,16 +140,28 @@ if __name__ == "__main__":
     # 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.
+    # see it deserialized instead of recompiled. One rule: ``x_initial`` is
+    # ``(B, n_x)`` — one extra leading axis — so it is batched; the terminal
+    # pin is omitted, so every element shares the default.
     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)
+    batched = export_problem.solve_batched(x_initial=x_initial_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)}")
+
+    # Algorithm knobs batch by the same rule, through the `algorithm` dict —
+    # the same names as the Problem constructor's algorithm config. ep_tr is
+    # a scalar field, so a (B,) vector sweeps the convergence tolerance per
+    # element — same artifact, no recompile, since the knobs are runtime
+    # inputs on the state pytree. The batch size is never passed: it is read
+    # off the array's leading axis (4 tolerances -> B = 4).
+    sweep = export_problem.solve_batched(
+        algorithm={"ep_tr": jnp.logspace(-4, -1, 4), "ep_vc": 1e-7}
+    )
+    print(f"  ep_tr sweep t_final:   {np.asarray(sweep.t_final).reshape(-1)}")
+    print(f"  ep_tr sweep converged: {np.asarray(sweep.converged)}")

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/examples/rocket/6DoF_pdg_batched_ic.py

@@ -2,7 +2,7 @@
 
 Runs ``N_BATCH`` SCP solves in parallel, each starting from a different random
 initial position drawn uniformly within the position bounds.  Per-batch ICs are
-passed through ``x0_stack``: position is declared ``Fix`` at the symbolic level
+passed through ``x_initial``: position is declared ``Fix`` at the symbolic level
 (so the solver consumes ``x_init_pin``), then each batch row overwrites the
 position slice with the sampled values.  This avoids the per-iteration
 ``(position == parameter).convex()`` constraint while still enforcing the full
@@ -110,7 +110,7 @@ mass.final = [ox.Maximize(1.5)]
 position = ox.State("position", shape=(3,))
 position.max = [10.0, 10.0, 10.0]
 position.min = [-10.0, -10.0, -10.0]
-# Fix so x0_stack pins are consumed; per-batch values overwrite these defaults.
+# Fix so x_initial pins are consumed; per-batch values overwrite these defaults.
 position.initial = [7.5, 4.5, 2.5]
 position.final = [0.0, ox.Free(0.0), ox.Free(0.0)]
 
@@ -193,7 +193,7 @@ constraint_exprs = []
 for st in states:
     constraint_exprs.extend([ox.ctcs(st <= st.max), ox.ctcs(st.min <= st)])
 
-# Initial position via x0_stack (Fix BC); terminal lateral via convex constraint.
+# Initial position via x_initial pins (Fix BC); terminal lateral via convex constraint.
 constraint_exprs.append((position[1:3] == final_position).convex().at([N - 1]))
 
 constraint_exprs.append(ox.ctcs(1.0 * (mass - m_dry) >= 0))
@@ -414,15 +414,14 @@ if __name__ == "__main__":
     print(f"Sampled {N_BATCH} initial positions  (alt ∈ [3, 9] m, lat ∈ [-7, 7] m)")
 
     # ── Compile + solve ───────────────────────────────────────────────────────
+    # x_initial is (B, n_x) — one extra leading axis — so it is batched; the
+    # terminal pin is omitted, so every element shares the default.
     default_init = np.asarray(problem.state.x_init_pin)
-    x0_stack = np.broadcast_to(default_init, (N_BATCH, default_init.shape[0])).copy()
-    x0_stack[:, POSITION_SLICE] = ic_batch
-    x0_stack = jnp.asarray(x0_stack)
-    xf_pin = np.asarray(problem.state.x_term_pin)
-    xf_stack = jnp.broadcast_to(xf_pin, (N_BATCH, xf_pin.shape[0]))
+    x_initial = np.broadcast_to(default_init, (N_BATCH, default_init.shape[0])).copy()
+    x_initial[:, POSITION_SLICE] = ic_batch
 
     print("Compiling and running solve_batched …")
-    results = problem.solve_batched(x0_stack=x0_stack, xf_stack=xf_stack)
+    results = problem.solve_batched(x_initial=jnp.asarray(x_initial))
 
     # ── Post-process: propagate all B solutions through nonlinear dynamics ────
     print("Post-processing (nonlinear propagation) …")

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/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.dev23'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev23')
+__version__ = version = '0.5.3.dev24'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev24')
 
-__commit_id__ = commit_id = 'g107a1d50b'
+__commit_id__ = commit_id = 'gba72d9b80'

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/openscvx/algorithms/__init__.py

@@ -53,6 +53,7 @@ from .base import (
     AlgorithmState,
     AutotuningBase,
     DiscretizationResult,
+    HyperParams,
     adaptive_state_code_to_str,
 )
 from .optimization_results import OptimizationResults
@@ -145,6 +146,7 @@ __all__ = [
     # PTR algorithm
     "PenalizedTrustRegion",
     "AutotuningBase",
+    "HyperParams",
     "AugmentedLagrangian",
     "AdaptiveProximalWeight",
     "ConstantProximalWeight",

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/openscvx/algorithms/autotuner/adaptive_proximal_weight.py

@@ -13,7 +13,7 @@ from pydantic import BaseModel, ConfigDict
 
 from openscvx.config import Config
 
-from ..base import AdaptiveStateCode, AutotuningBase
+from ..base import AdaptiveStateCode, AutotuningBase, HyperParams
 from .augmented_lagrangian import AugmentedLagrangian
 
 if TYPE_CHECKING:
@@ -22,6 +22,12 @@ if TYPE_CHECKING:
     from ..base import AlgorithmState, CandidateIterate
 
 
+class AdaptiveProximalWeightHyper(HyperParams):
+    """Declared hyperparameters for :class:`AdaptiveProximalWeight`."""
+
+    lam_cost_drop: int = -1
+
+
 class AdaptiveProximalWeight(AutotuningBase):
     """PTR-style proximal adaptation with fixed virtual-penalty weights.
 
@@ -54,7 +60,7 @@ class AdaptiveProximalWeight(AutotuningBase):
         self.eta_2 = eta_2
         self.lam_prox_min = lam_prox_min
         self.lam_prox_max = lam_prox_max
-        self.lam_cost_drop = lam_cost_drop
+        self.hyper = AdaptiveProximalWeightHyper(lam_cost_drop=lam_cost_drop)
         self.lam_cost_relax = lam_cost_relax
 
     def update_weights(
@@ -85,7 +91,7 @@ class AdaptiveProximalWeight(AutotuningBase):
         J_nonlin = nonlin_cost + nonlin_pen + nodal_pen
 
         lam_cost_next = jnp.where(
-            state.k > self.lam_cost_drop,
+            state.k > state.hyper.lam_cost_drop,
             state.lam_cost * self.lam_cost_relax,
             state.lam_cost_init,
         )

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/openscvx/algorithms/autotuner/augmented_lagrangian.py

@@ -19,7 +19,7 @@ from openscvx.utils.printing import (
     color_adaptive_state,
 )
 
-from ..base import AdaptiveStateCode, AutotuningBase
+from ..base import AdaptiveStateCode, AutotuningBase, HyperParams
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
@@ -27,6 +27,14 @@ if TYPE_CHECKING:
     from ..base import AlgorithmState, CandidateIterate
 
 
+class AugmentedLagrangianHyper(HyperParams):
+    """Declared hyperparameters for :class:`AugmentedLagrangian`."""
+
+    rho_init: float = 1.0
+    rho_max: float = 1e2
+    lam_cost_drop: int = -1
+
+
 class AugmentedLagrangian(AutotuningBase):
     """Augmented Lagrangian autotuner.
 
@@ -80,7 +88,12 @@ class AugmentedLagrangian(AutotuningBase):
         """Initialize Augmented Lagrangian autotuning parameters.
 
         All parameters have defaults and can be modified after instantiation
-        via attribute access (e.g., ``autotuner.rho_max = 1e7``).
+        via attribute access (e.g., ``autotuner.lam_prox_max = 1e6``); the
+        declared hyperparameters (``rho_init`` / ``rho_max`` /
+        ``lam_cost_drop``) live on the frozen ``hyper`` container instead
+        (``autotuner.hyper = dataclasses.replace(autotuner.hyper,
+        rho_max=1e7)``) and are also per-solve overrides — see
+        :class:`AutotuningBase`.
 
         Args:
             rho_init: Initial penalty parameter for constraints. Defaults to 1.0.
@@ -106,8 +119,11 @@ class AugmentedLagrangian(AutotuningBase):
             lam_cost_relax: Factor applied to lam_cost after lam_cost_drop.
                 Defaults to 1.0.
         """
-        self.rho_init = rho_init
-        self.rho_max = rho_max
+        self.hyper = AugmentedLagrangianHyper(
+            rho_init=rho_init,
+            rho_max=rho_max,
+            lam_cost_drop=lam_cost_drop,
+        )
         self.gamma_1 = gamma_1
         self.gamma_2 = gamma_2
         self.eta_0 = eta_0
@@ -118,7 +134,6 @@ class AugmentedLagrangian(AutotuningBase):
         self.lam_vc_max = lam_vc_max
         self.lam_prox_min = lam_prox_min
         self.lam_prox_max = lam_prox_max
-        self.lam_cost_drop = lam_cost_drop
         self.lam_cost_relax = lam_cost_relax
 
     # -----------------------------------------------------------------------
@@ -244,12 +259,14 @@ class AugmentedLagrangian(AutotuningBase):
         )
         J_nonlin = nonlin_cost + nonlin_pen + nodal_pen
 
-        # Cost relaxation: when state.k > lam_cost_drop, scale state.lam_cost;
-        # otherwise reset to the algorithm's initial weight (carried on the
-        # pytree as state.lam_cost_init, broadcast at from_settings()). Scalar
+        # Cost relaxation: when state.k > hyper.lam_cost_drop, scale
+        # state.lam_cost; otherwise reset to the algorithm's initial weight
+        # (carried on the pytree as state.lam_cost_init, broadcast at
+        # from_settings()). Both constants ride the pytree so per-solve
+        # overrides and vmap sweeps reach the traced body. Scalar
         # lam_cost_relax preserves the user-specified per-state weight ratios.
         lam_cost_next = jnp.where(
-            state.k > self.lam_cost_drop,
+            state.k > state.hyper.lam_cost_drop,
             state.lam_cost * self.lam_cost_relax,
             state.lam_cost_init,
         )

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/openscvx/algorithms/autotuner/constant_proximal_weight.py

@@ -11,7 +11,7 @@ from pydantic import BaseModel, ConfigDict
 
 from openscvx.config import Config
 
-from ..base import AdaptiveStateCode, AutotuningBase
+from ..base import AdaptiveStateCode, AutotuningBase, HyperParams
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
@@ -19,6 +19,12 @@ if TYPE_CHECKING:
     from ..base import AlgorithmState, CandidateIterate
 
 
+class ConstantProximalWeightHyper(HyperParams):
+    """Declared hyperparameters for :class:`ConstantProximalWeight`."""
+
+    lam_cost_drop: int = -1
+
+
 class ConstantProximalWeight(AutotuningBase):
     """Hold ``lam_prox`` constant; relax ``lam_cost`` after ``lam_cost_drop``.
 
@@ -38,7 +44,7 @@ class ConstantProximalWeight(AutotuningBase):
         lam_cost_drop: int = -1,
         lam_cost_relax: float = 1.0,
     ):
-        self.lam_cost_drop = lam_cost_drop
+        self.hyper = ConstantProximalWeightHyper(lam_cost_drop=lam_cost_drop)
         self.lam_cost_relax = lam_cost_relax
 
     def update_weights(
@@ -54,7 +60,7 @@ class ConstantProximalWeight(AutotuningBase):
         Pure functional update — see class docstring.
         """
         lam_cost_next = jnp.where(
-            state.k > self.lam_cost_drop,
+            state.k > state.hyper.lam_cost_drop,
             state.lam_cost * self.lam_cost_relax,
             state.lam_cost_init,
         )

{openscvx-0.5.3.dev23 → openscvx-0.5.3.dev24}/openscvx/algorithms/autotuner/ramp_proximal_weight.py

@@ -11,7 +11,7 @@ from pydantic import BaseModel, ConfigDict
 
 from openscvx.config import Config
 
-from ..base import AdaptiveStateCode, AutotuningBase
+from ..base import AdaptiveStateCode, AutotuningBase, HyperParams
 
 if TYPE_CHECKING:
     from openscvx.lowered import LoweredJaxConstraints
@@ -19,6 +19,12 @@ if TYPE_CHECKING:
     from ..base import AlgorithmState, CandidateIterate
 
 
+class RampProximalWeightHyper(HyperParams):
+    """Declared hyperparameters for :class:`RampProximalWeight`."""
+
+    lam_cost_drop: int = -1
+
+
 class RampProximalWeight(AutotuningBase):
     """Ramp ``lam_prox`` toward ``lam_prox_max`` then hold.
 
@@ -40,7 +46,7 @@ class RampProximalWeight(AutotuningBase):
     ):
         self.ramp_factor = ramp_factor
         self.lam_prox_max = lam_prox_max
-        self.lam_cost_drop = lam_cost_drop
+        self.hyper = RampProximalWeightHyper(lam_cost_drop=lam_cost_drop)
         self.lam_cost_relax = lam_cost_relax
 
     def update_weights(
@@ -56,7 +62,7 @@ class RampProximalWeight(AutotuningBase):
         Pure functional update — see class docstring.
         """
         lam_cost_next = jnp.where(
-            state.k > self.lam_cost_drop,
+            state.k > state.hyper.lam_cost_drop,
             state.lam_cost * self.lam_cost_relax,
             state.lam_cost_init,
         )