openscvx 0.5.3.dev36__tar.gz → 0.5.3.dev38__tar.gz

{openscvx-0.5.3.dev36/openscvx.egg-info → openscvx-0.5.3.dev38}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev36
+Version: 0.5.3.dev38
 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.dev36 → openscvx-0.5.3.dev38}/examples/spacecraft/relative_loitering.py

@@ -353,7 +353,6 @@ def build_relative_loitering_problem(
         time=time,
         constraints=constraints,
         N=n_nodes,
-        licq_min=0.0,
         licq_max=1e-8,
         discretizer={
             "ode_solver": "Dopri8",

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/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.dev36'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev36')
+__version__ = version = '0.5.3.dev38'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev38')
 
-__commit_id__ = commit_id = 'g1eddc1a14'
+__commit_id__ = commit_id = 'g235acbb25'

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/problem.py

@@ -340,8 +340,7 @@ class Problem:
         dynamics_prop: Optional[dict] = None,
         states_prop: Optional[List[State]] = None,
         algebraic_prop: Optional[dict] = None,
-        licq_min: Union[float, Dict[int, float]] = 0.0,
-        licq_max: Union[float, Dict[int, float]] = 1e-4,
+        licq_max: float = 1e-4,
         algorithm: Optional[Union[Algorithm, dict]] = None,
         discretizer: Optional[Union[Discretizer, dict]] = None,
         solver: Optional[Union[ConvexSolver, dict]] = None,
@@ -373,12 +372,11 @@ class Problem:
                 Only specify additional states beyond optimization states. Used with dynamics_prop.
             algebraic_prop (dict, optional): Dictionary mapping names to symbolic expressions
                 for outputs evaluated (not integrated) during propagation.
-            licq_min: Minimum LICQ constraint value. Defaults to 0.0.
-                Either a scalar (applied to all CTCS groups) or a dict
-                mapping CTCS group ``idx`` to per-group bounds.
-            licq_max: Maximum LICQ constraint value. Defaults to 1e-4.
-                Either a scalar (applied to all CTCS groups) or a dict
-                mapping CTCS group ``idx`` to per-group bounds.
+            licq_max: Problem-wide upper bound for CTCS augmented states
+                (default: 1e-4). Serves as the default for all constraint
+                groups; a ``licq_max`` set on an individual constraint (via
+                ``ctcs(...)`` or ``.over(...)``) overrides it for that
+                constraint's augmented state.
             algorithm: SCP algorithm configuration. Accepts:
 
                 - ``None`` — uses ``PenalizedTrustRegion()`` with defaults.
@@ -539,7 +537,6 @@ class Problem:
             controls=controls,
             N=N,
             time=time,
-            licq_min=licq_min,
             licq_max=licq_max,
             dynamics_prop_extra=dynamics_prop,
             states_prop_extra=states_prop,

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/symbolic/augmentation.py

@@ -54,7 +54,7 @@ Example:
         # controls_aug includes original controls + time dilation
 """
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -448,8 +448,7 @@ def augment_dynamics_with_ctcs(
     N: int,
     xdelta: Optional[Expr] = None,
     *,
-    licq_min: Union[float, Dict[int, float]] = 0.0,
-    licq_max: Union[float, Dict[int, float]] = 1e-4,
+    licq_max: float = 1e-4,
 ) -> Tuple[Expr, Optional[Expr], List[State], List[Control]]:
     """Augment dynamics with continuous-time constraint satisfaction states.
 
@@ -467,18 +466,20 @@ def augment_dynamics_with_ctcs(
         aug_dot = penalty(g(x, u))  # For each constraint group
         time_dot = 1.0
 
+    Each augmented state is bounded to [0, licq_max], where the group's bound
+    is taken from its member constraints' ``licq_max`` attribute (constraints in
+    a group must agree) and falls back to the problem-wide ``licq_max`` when no
+    member specifies one.
+
     Args:
         xdot: Original dynamics expression for states
         states: List of state variables (must include a state named "time")
         controls: List of control variables
         constraints_ctcs: List of CTCS constraints (should be sorted and grouped)
         N: Number of discretization nodes
-        licq_min: Minimum bound for augmented states (default: 0.0).
-            Either a scalar (applied to all groups) or a dict mapping
-            CTCS group ``idx`` to per-group bounds.
-        licq_max: Maximum bound for augmented states (default: 1e-4).
-            Either a scalar (applied to all groups) or a dict mapping
-            CTCS group ``idx`` to per-group bounds.
+        licq_max: Problem-wide default upper bound for augmented states
+            (default: 1e-4). Per-constraint ``licq_max`` values override it
+            for their group.
 
     Returns:
         Tuple of:
@@ -489,6 +490,8 @@ def augment_dynamics_with_ctcs(
 
     Raises:
         ValueError: If no state named "time" is found in the states list
+        ValueError: If constraints sharing an augmented state specify
+            conflicting ``licq_max`` values
 
     Example:
         Augment dynamics with CTCS penalty states:
@@ -509,6 +512,9 @@ def augment_dynamics_with_ctcs(
         states_aug includes x, time, and _ctcs_aug_0,
         controls_aug includes u and _time_dilation
     """
+    if licq_max <= 0:
+        raise ValueError(f"licq_max must be a positive number, got {licq_max!r}")
+
     # Save if discrete dynamics is null
     discrete_dynamics_set = xdelta is not None
 
@@ -535,24 +541,6 @@ def augment_dynamics_with_ctcs(
                 penalty_groups[ctcs.idx] = []
             penalty_groups[ctcs.idx].append(ctcs)
 
-        # Validate dict-valued licq_min / licq_max keys against actual groups.
-        valid_ids = set(penalty_groups.keys())
-        for name, val in [("licq_min", licq_min), ("licq_max", licq_max)]:
-            if isinstance(val, dict):
-                unknown = set(val.keys()) - valid_ids
-                if unknown:
-                    raise ValueError(
-                        f"{name} dict contains unknown CTCS group idx: {unknown}. "
-                        f"Valid idx values: {sorted(valid_ids)}"
-                    )
-                missing = valid_ids - set(val.keys())
-                if missing:
-                    raise ValueError(
-                        f"{name} dict is missing entries for CTCS group idx: "
-                        f"{sorted(missing)}. When using a dict, all groups must "
-                        f"have an entry. Valid idx values: {sorted(valid_ids)}"
-                    )
-
         # Create augmented state expressions for each group and corresponding
         # augmented state variables. For discrete dynamics, preserve these states
         # by identity mapping across impulsive updates.
@@ -566,17 +554,25 @@ def augment_dynamics_with_ctcs(
                 augmented_state_expr = Add(*penalty_terms)
             augmented_state_exprs.append(augmented_state_expr)
 
-            # Resolve per-group LICQ bounds
-            idx_min = licq_min[idx] if isinstance(licq_min, dict) else licq_min
-            idx_max = licq_max[idx] if isinstance(licq_max, dict) else licq_max
+            # Resolve the group's LICQ bound from its member constraints:
+            # explicit values must agree; otherwise fall back to the
+            # problem-wide default.
+            explicit = {c.licq_max for c in penalty_terms if c.licq_max is not None}
+            if len(explicit) > 1:
+                raise ValueError(
+                    f"CTCS constraints sharing augmented state idx={idx} specify "
+                    f"conflicting licq_max values: {sorted(explicit)}. Constraints "
+                    f"in the same group must agree on licq_max; set the same value "
+                    f"on each, or separate them into different groups via idx."
+                )
+            group_licq_max = explicit.pop() if explicit else licq_max
 
             aug_var = State(f"_ctcs_aug_{idx}", shape=(1,))
-            aug_var.initial = np.array([idx_min])  # Set initial to respect bounds
+            aug_var.initial = np.array([0.0])
             aug_var.final = [("free", 0)]
-            aug_var.min = np.array([idx_min])
-            aug_var.max = np.array([idx_max])
-            # Set guess to idx_min as well
-            aug_var.guess = np.full([N, 1], idx_min)  # N x num augmented states
+            aug_var.min = np.array([0.0])
+            aug_var.max = np.array([group_licq_max])
+            aug_var.guess = np.zeros([N, 1])  # N x num augmented states
             states_augmented.append(aug_var)
             augmented_state_exprs_discrete.append(aug_var)
 

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/symbolic/builder.py

@@ -29,7 +29,7 @@ Pipeline stages:
 See `preprocess_symbolic_problem()` for the main entry point.
 """
 
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 
@@ -71,8 +71,7 @@ def preprocess_symbolic_problem(
     controls: List[Control],
     N: int,
     time: Time,
-    licq_min: Union[float, Dict[int, float]] = 0.0,
-    licq_max: Union[float, Dict[int, float]] = 1e-4,
+    licq_max: float = 1e-4,
     dynamics_prop_extra: dict = None,
     dynamics_discrete: dict = None,
     states_prop_extra: List[State] = None,
@@ -106,12 +105,9 @@ def preprocess_symbolic_problem(
         controls: List of user-defined Control objects (should NOT include time dilation)
         N: Number of discretization nodes in the trajectory
         time: Time configuration object specifying time bounds and constraints
-        licq_min: Minimum bound for CTCS augmented states (default: 0.0).
-            Either a scalar (applied to all groups) or a dict mapping
-            CTCS group ``idx`` to per-group bounds.
-        licq_max: Maximum bound for CTCS augmented states (default: 1e-4).
-            Either a scalar (applied to all groups) or a dict mapping
-            CTCS group ``idx`` to per-group bounds.
+        licq_max: Problem-wide default upper bound for CTCS augmented states
+            (default: 1e-4). Per-constraint ``licq_max`` values override it
+            for their group.
         dynamics_prop_extra: Optional dictionary of additional dynamics for propagation-only
             states (default: None)
         states_prop_extra: Optional list of additional State objects for propagation only
@@ -350,7 +346,6 @@ def preprocess_symbolic_problem(
         constraints.ctcs,
         N,
         xdelta=dynamics_discrete_concat,
-        licq_min=licq_min,
         licq_max=licq_max,
     )
 

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/symbolic/expr/constraint.py

@@ -127,6 +127,7 @@ class Constraint(Expr):
         penalty: str = "squared_relu",
         idx: Optional[int] = None,
         check_nodally: bool = False,
+        licq_max: Optional[float] = None,
     ) -> "CTCS":
         """Apply this constraint over a continuous interval using CTCS.
 
@@ -135,11 +136,22 @@ class Constraint(Expr):
             penalty: Penalty function type ("squared_relu", "huber", "smooth_relu")
             idx: Optional grouping index for multiple augmented states
             check_nodally: Whether to also enforce this constraint nodally
+            licq_max: Optional upper bound on the augmented state that accumulates
+                this constraint's violation penalty. Constraints sharing an
+                augmented state must agree on the value; groups without one use
+                the problem-wide ``licq_max`` (a ``Problem`` argument).
 
         Returns:
             CTCS constraint wrapping this constraint with interval specification
         """
-        return CTCS(self, penalty=penalty, nodes=interval, idx=idx, check_nodally=check_nodally)
+        return CTCS(
+            self,
+            penalty=penalty,
+            nodes=interval,
+            idx=idx,
+            check_nodally=check_nodally,
+            licq_max=licq_max,
+        )
 
     def convex(self) -> "Constraint":
         """Mark this constraint as convex for CVXPy lowering.
@@ -609,6 +621,11 @@ class CTCS(Expr):
             an augmented state. If None, auto-assigned based on node intervals.
         check_nodally: Whether to also enforce the constraint at discrete nodes for
             additional numerical robustness (creates both continuous and nodal constraints)
+        licq_max: Optional upper bound on the augmented state accumulating this
+            constraint's violation penalty. Smaller values enforce the constraint
+            more strictly between nodes. Constraints sharing an augmented state
+            must agree on the value; groups without one use the problem-wide
+            ``licq_max`` (a ``Problem`` argument).
 
     Example:
         Single augmented state (default behavior - same node interval):
@@ -642,6 +659,7 @@ class CTCS(Expr):
         nodes: Optional[Tuple[int, int]] = None,
         idx: Optional[int] = None,
         check_nodally: bool = False,
+        licq_max: Optional[float] = None,
     ):
         """Initialize a CTCS constraint.
 
@@ -662,11 +680,16 @@ class CTCS(Expr):
             check_nodally: If True, also enforce the constraint at discrete nodes for
                 numerical stability (creates both continuous and nodal constraints).
                 Defaults to False.
+            licq_max: Optional upper bound on the augmented state accumulating this
+                constraint's violation penalty. Constraints sharing an augmented
+                state must agree on the value; groups without one use the
+                problem-wide ``licq_max`` (a ``Problem`` argument).
 
         Raises:
             TypeError: If constraint is not a Constraint instance
             ValueError: If nodes is not None or a 2-tuple of integers
             ValueError: If nodes[0] >= nodes[1] (invalid interval)
+            ValueError: If licq_max is not a positive number
         """
         if not isinstance(constraint, Constraint):
             raise TypeError("CTCS must wrap a Constraint")
@@ -683,12 +706,19 @@ class CTCS(Expr):
             if nodes[0] >= nodes[1]:
                 raise ValueError("CTCS node range must have start < end")
 
+        if licq_max is not None:
+            if not isinstance(licq_max, (int, float)) or licq_max <= 0:
+                raise ValueError(f"licq_max must be a positive number, got {licq_max!r}")
+            licq_max = float(licq_max)
+
         self.constraint = constraint
         self.penalty = penalty
         self.nodes = nodes  # (start, end) node range or None for all nodes
         self.idx = idx  # Optional grouping index for multiple augmented states
         # Whether to also enforce this constraint nodally for numerical stability
         self.check_nodally = check_nodally
+        # Upper bound on the augmented state; None defers to the group default
+        self.licq_max = licq_max
 
     def children(self) -> List[Expr]:
         """Return the wrapped constraint as the only child.
@@ -711,6 +741,7 @@ class CTCS(Expr):
             nodes=self.nodes,
             idx=self.idx,
             check_nodally=self.check_nodally,
+            licq_max=self.licq_max,
         )
 
     def check_shape(self) -> Tuple[int, ...]:
@@ -767,6 +798,11 @@ class CTCS(Expr):
             hasher.update(b"None")
         # Hash check_nodally
         hasher.update(b"1" if self.check_nodally else b"0")
+        # Hash licq_max
+        if self.licq_max is not None:
+            hasher.update(struct.pack(">d", self.licq_max))
+        else:
+            hasher.update(b"None")
         # Hash the wrapped constraint
         self.constraint._hash_into(hasher)
 
@@ -794,6 +830,7 @@ class CTCS(Expr):
             nodes=interval,
             idx=self.idx,
             check_nodally=self.check_nodally,
+            licq_max=self.licq_max,
         )
 
     def __repr__(self) -> str:
@@ -809,6 +846,8 @@ class CTCS(Expr):
             parts.append(f"idx={self.idx}")
         if self.check_nodally:
             parts.append(f"check_nodally={self.check_nodally}")
+        if self.licq_max is not None:
+            parts.append(f"licq_max={self.licq_max}")
         return f"CTCS({', '.join(parts)})"
 
     def penalty_expr(self) -> Expr:
@@ -861,6 +900,7 @@ def ctcs(
     nodes: Optional[Tuple[int, int]] = None,
     idx: Optional[int] = None,
     check_nodally: bool = False,
+    licq_max: Optional[float] = None,
 ) -> CTCS:
     """Helper function to create CTCS (Continuous-Time Constraint Satisfaction) constraints.
 
@@ -876,6 +916,10 @@ def ctcs(
         idx: Optional grouping index for multiple augmented states
         check_nodally: Whether to also enforce constraint at discrete nodes.
             Defaults to False.
+        licq_max: Optional upper bound on the augmented state accumulating this
+            constraint's violation penalty. Constraints sharing an augmented
+            state must agree on the value; groups without one use the
+            problem-wide ``licq_max`` (a ``Problem`` argument).
 
     Returns:
         CTCS: A CTCS constraint wrapping the input constraint
@@ -899,4 +943,4 @@ def ctcs(
 
             altitude_constraint = (altitude >= 10).over((0, 100), penalty="huber")
     """
-    return CTCS(constraint, penalty, nodes, idx, check_nodally)
+    return CTCS(constraint, penalty, nodes, idx, check_nodally, licq_max)

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/symbolic/expr/expr.py

@@ -63,8 +63,10 @@ Example:
         canonical = expr.canonicalize()  # Simplifies to: x + x
 """
 
+import functools
 import hashlib
 import struct
+import threading
 from typing import TYPE_CHECKING, Callable, List, Tuple, Union
 
 if TYPE_CHECKING:
@@ -72,6 +74,40 @@ if TYPE_CHECKING:
 
 import numpy as np
 
+# Per-thread memo for canonicalize(). The memo canonicalizes each node once and
+# returns the same object to every parent, preserving sharing.
+_canon_state = threading.local()
+
+
+def _memoized_canonicalize(fn: Callable) -> Callable:
+    """Wrap a subclass ``canonicalize`` so shared subexpressions are reused.
+
+    Keyed by ``id(self)``; the node is held alive in the memo for the duration of
+    the pass so its id cannot be recycled, with an identity guard on lookup. The
+    top-level call owns the memo and clears it on exit, so each independent
+    canonicalize pass starts fresh.
+    """
+
+    @functools.wraps(fn)
+    def wrapper(self):
+        memo = getattr(_canon_state, "memo", None)
+        owns = memo is None
+        if owns:
+            memo = {}
+            _canon_state.memo = memo
+        try:
+            entry = memo.get(id(self))
+            if entry is not None and entry[0] is self:
+                return entry[1]
+            result = fn(self)
+            memo[id(self)] = (self, result)
+            return result
+        finally:
+            if owns:
+                _canon_state.memo = None
+
+    return wrapper
+
 
 class Expr:
     """Base class for symbolic expressions in optimization problems.
@@ -101,6 +137,13 @@ class Expr:
     # Give Expr objects higher priority than numpy arrays in operations
     __array_priority__ = 1000
 
+    def __init_subclass__(cls, **kwargs):
+        # Wrap each subclass's own canonicalize() with identity memoization so a
+        # shared DAG node is canonicalized once rather than re-expanded per parent.
+        super().__init_subclass__(**kwargs)
+        if "canonicalize" in cls.__dict__:
+            cls.canonicalize = _memoized_canonicalize(cls.__dict__["canonicalize"])
+
     def __le__(self, other):
         from .constraint import Inequality
 

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/symbolic/expr/stl.py

@@ -310,6 +310,7 @@ class STLExpr(Expr):
         penalty: str = "smooth_relu",
         idx: Optional[int] = None,
         check_nodally: bool = False,
+        licq_max: Optional[float] = None,
         interval_type: IntervalKind = "nodes",
     ) -> "CTCS":
         """Apply this STL expression over an enforcement interval using CTCS.
@@ -323,6 +324,10 @@ class STLExpr(Expr):
             penalty: Penalty function type for CTCS
             idx: Optional grouping index for multiple augmented states
             check_nodally: Whether to also enforce at discrete nodes
+            licq_max: Optional upper bound on the augmented state accumulating
+                this constraint's violation penalty. Constraints sharing an
+                augmented state must agree on the value; groups without one use
+                the problem-wide ``licq_max`` (a ``Problem`` argument).
 
         Returns:
             Continuous-time constraint satisfaction wrapper
@@ -359,6 +364,7 @@ class STLExpr(Expr):
             nodes=(coerced.start, coerced.end),
             idx=idx,
             check_nodally=check_nodally,
+            licq_max=licq_max,
         )
 
     def at(self, nodes: Union[list, tuple]) -> "NodalConstraint":
@@ -863,6 +869,7 @@ class Always(_TemporalSTLExpr):
         penalty: str = "smooth_relu",
         idx: Optional[int] = None,
         check_nodally: bool = False,
+        licq_max: Optional[float] = None,
         interval_type: IntervalKind = "nodes",
     ) -> "CTCS":
         """Convert this ``Always`` to a ``CTCS`` constraint.
@@ -876,6 +883,10 @@ class Always(_TemporalSTLExpr):
             penalty: CTCS penalty function name.
             idx: Optional grouping index for multiple augmented states.
             check_nodally: Whether to additionally enforce at discrete nodes.
+            licq_max: Optional upper bound on the augmented state accumulating
+                this constraint's violation penalty. Constraints sharing an
+                augmented state must agree on the value; groups without one use
+                the problem-wide ``licq_max`` (a ``Problem`` argument).
 
         Returns:
             A ``CTCS`` constraint enforcing the inner predicate over the
@@ -903,7 +914,7 @@ class Always(_TemporalSTLExpr):
 
         if isinstance(self.predicate, STLExpr):
             return self.predicate.over(
-                chosen, penalty=penalty, idx=idx, check_nodally=check_nodally
+                chosen, penalty=penalty, idx=idx, check_nodally=check_nodally, licq_max=licq_max
             )
 
         from .constraint import CTCS
@@ -914,6 +925,7 @@ class Always(_TemporalSTLExpr):
             nodes=(chosen.start, chosen.end),
             idx=idx,
             check_nodally=check_nodally,
+            licq_max=licq_max,
         )
 
     def __repr__(self) -> str:

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev38}/openscvx/symbolic/expr/stljax.py

@@ -59,6 +59,7 @@ class STLJaxExpr(Expr):
         penalty: str = "squared_relu",
         idx: Optional[int] = None,
         check_nodally: bool = False,
+        licq_max: Optional[float] = None,
     ) -> "CTCS":
         """Apply this STL expression over a continuous interval using CTCS.
 
@@ -70,6 +71,10 @@ class STLJaxExpr(Expr):
             penalty: Penalty function type for CTCS
             idx: Optional grouping index for multiple augmented states
             check_nodally: Whether to also enforce at discrete nodes
+            licq_max: Optional upper bound on the augmented state accumulating
+                this constraint's violation penalty. Constraints sharing an
+                augmented state must agree on the value; groups without one use
+                the problem-wide ``licq_max`` (a ``Problem`` argument).
 
         Returns:
             Continuous-time constraint satisfaction wrapper
@@ -89,7 +94,12 @@ class STLJaxExpr(Expr):
         constraint = Inequality(Neg(self), Constant(np.array(0.0)))
 
         return CTCS(
-            constraint, penalty=penalty, nodes=interval, idx=idx, check_nodally=check_nodally
+            constraint,
+            penalty=penalty,
+            nodes=interval,
+            idx=idx,
+            check_nodally=check_nodally,
+            licq_max=licq_max,
         )
 
     def at(self, nodes: Union[list, tuple]) -> "NodalConstraint":

openscvx-0.5.3.dev38/openscvx/symbolic/lowerers/jax/_lowerer.py

@@ -0,0 +1,122 @@
+"""JaxLowerer class definition.
+
+The visitor methods that populate the registry live in sibling modules
+(``arithmetic``, ``math``, ``linalg``, etc.) and are registered via
+``@visitor`` at import time.
+"""
+
+from typing import Callable
+
+import jax
+
+from openscvx.symbolic.expr import Expr
+from openscvx.symbolic.lowerers.jax._registry import dispatch
+
+_MISSING = object()
+
+# Value-memoization is only valid within a single JAX trace. Visitors bracket
+# their nested trace with :func:`pause_memo` / :func:`resume_memo` so the wrapped
+# closures recompute instead of returning a cross-trace value.
+_memo_paused = [0]
+
+
+def pause_memo() -> None:
+    """Disable value-memoization for the duration of a nested sub-trace."""
+    _memo_paused[0] += 1
+
+
+def resume_memo() -> None:
+    """Re-enable value-memoization after a nested sub-trace finishes tracing."""
+    _memo_paused[0] -= 1
+
+
+def _memoize_call(fn: Callable) -> Callable:
+    """Cache a lowered closure's result so it emits its subgraph once per trace.
+
+    Within a trace every node receives the same argument objects, so caching on
+    argument identity (``is``) lets a shared subexpression emit its subgraph once
+    instead of once per consumer. The cache only applies while tracing (at least
+    one positional argument is a JAX tracer); eager calls, keyword-argument
+    calls, and paused sections (nested sub-traces, see module note) bypass it.
+    """
+    last_args = None
+    last_val = _MISSING
+
+    def wrapped(*args, **kwargs):
+        nonlocal last_args, last_val
+        if kwargs or _memo_paused[0] or not any(isinstance(a, jax.core.Tracer) for a in args):
+            return fn(*args, **kwargs)
+        if (
+            last_val is not _MISSING
+            and last_args is not None
+            and len(last_args) == len(args)
+            and all(a is b for a, b in zip(last_args, args))
+        ):
+            return last_val
+        last_val = fn(*args)
+        last_args = args
+        return last_val
+
+    return wrapped
+
+
+class JaxLowerer:
+    """JAX backend for lowering symbolic expressions to executable functions.
+
+    This class implements the visitor pattern for converting symbolic expression
+    AST nodes to JAX functions. Each expression type has a corresponding visitor
+    function decorated with @visitor that handles the lowering logic.
+
+    The lowering process is recursive: each visitor lowers its child expressions
+    first, then composes them into a JAX operation. All lowered functions have
+    a standardized signature (x, u, node, params) -> result.
+
+    Shared subexpressions are lowered once: :meth:`lower` caches the closure for
+    each AST node by object identity, and each closure is wrapped so it emits its
+    subgraph only once per trace (see :func:`_memoize_call`).
+
+    Example:
+        Set up the JaxLowerer and lower an expression to a JAX function::
+
+            lowerer = JaxLowerer()
+            expr = ox.Norm(x)**2 + 0.1 * ox.Norm(u)**2
+            f = lowerer.lower(expr)
+            result = f(x_val, u_val, node=0, params={})
+
+    Note:
+        A fresh lowerer is created per lowering pass; its node cache is keyed by
+        ``id(expr)`` and is valid only while the expression tree is alive.
+    """
+
+    def __init__(self) -> None:
+        # Maps id(expr) -> (expr, lowered closure) so each node is lowered once,
+        # no matter how many parents share it. Storing the expr itself keeps it
+        # alive: if it were garbage-collected, Python could reuse its id for a
+        # new node and the cache would return the wrong closure. The identity
+        # check on lookup guards against exactly that.
+        self._node_cache: dict[int, tuple[Expr, Callable]] = {}
+
+    def lower(self, expr: Expr) -> Callable:
+        """Lower a symbolic expression to a JAX function.
+
+        Main entry point for lowering. Delegates to dispatch() which looks up
+        the appropriate visitor method based on the expression type. Results are
+        memoized by node identity so shared subexpressions are lowered once.
+
+        Args:
+            expr: Symbolic expression to lower (any Expr subclass)
+
+        Returns:
+            JAX function with signature (x, u, node, params) -> result
+
+        Raises:
+            NotImplementedError: If no visitor exists for the expression type
+            ValueError: If the expression is malformed (e.g., State without slice)
+        """
+        eid = id(expr)
+        cached = self._node_cache.get(eid)
+        if cached is not None and cached[0] is expr:
+            return cached[1]
+        fn = _memoize_call(dispatch(self, expr))
+        self._node_cache[eid] = (expr, fn)
+        return fn