openscvx 0.5.3.dev36__tar.gz → 0.5.3.dev37__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev36
+Version: 0.5.3.dev37
 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.dev37}/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.dev37'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev37')
 
-__commit_id__ = commit_id = 'g1eddc1a14'
+__commit_id__ = commit_id = 'gdb86fef25'

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev37}/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.dev37/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

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev37}/openscvx/symbolic/lowerers/jax/logic.py

@@ -8,6 +8,7 @@ from jax.lax import cond
 
 # Expression types to handle — uncomment as you paste visitors:
 from openscvx.symbolic.expr.logic import All, Any, Cond
+from openscvx.symbolic.lowerers.jax._lowerer import pause_memo, resume_memo
 from openscvx.symbolic.lowerers.jax._registry import visitor  # noqa: F401
 
 # Module-level default dtype for conditional branches
@@ -189,11 +190,20 @@ def _visit_cond(lowerer, node: Cond):
         def _false_branch(_):
             return jnp.asarray(false_fn(x, u, node_arg, params), dtype=default_dtype)
 
-        return cond(
-            pred_bool,
-            _true_branch,
-            _false_branch,
-            operand=None,
-        )
+        # lax.cond traces each branch separately, but both branch closures see
+        # the same captured (x, u, node, params). Pause the memo so shared
+        # subexpressions are recomputed inside each branch: a value cached in
+        # one trace is not valid in another.
+        pause_memo()
+        try:
+            result = cond(
+                pred_bool,
+                _true_branch,
+                _false_branch,
+                operand=None,
+            )
+        finally:
+            resume_memo()
+        return result
 
     return cond_fn

{openscvx-0.5.3.dev36 → openscvx-0.5.3.dev37}/openscvx/symbolic/lowerers/jax/vmap.py

@@ -8,6 +8,7 @@ import jax.numpy as jnp
 
 # Expression types to handle — uncomment as you paste visitors:
 from openscvx.symbolic.expr.vmap import Vmap, _Placeholder
+from openscvx.symbolic.lowerers.jax._lowerer import pause_memo, resume_memo
 from openscvx.symbolic.lowerers.jax._registry import visitor  # noqa: F401
 
 
@@ -129,9 +130,14 @@ def _visit_vmap(lowerer, node: Vmap):
                     new_params[key] = v
                 return inner_fn(x, u, node_idx, new_params)
 
-            # vmap over all batch arguments along the same axis
+            # Nested sub-trace reusing captured (x, u, node_idx): pause
+            # value-memo so no cached tracer crosses the trace boundary.
             in_axes = tuple(axis for _ in range(num_batches))
-            return jax.vmap(inner, in_axes=in_axes)(*data_arrays)
+            pause_memo()
+            try:
+                return jax.vmap(inner, in_axes=in_axes)(*data_arrays)
+            finally:
+                resume_memo()
 
     else:
         # All Constants: bake all data into closure at lowering time
@@ -145,8 +151,12 @@ def _visit_vmap(lowerer, node: Vmap):
                     new_params[key] = v
                 return inner_fn(x, u, node_idx, new_params)
 
-            # vmap over all batch arguments along the same axis
+            # Nested sub-trace: pause value-memo (see runtime branch above).
             in_axes = tuple(axis for _ in range(num_batches))
-            return jax.vmap(inner, in_axes=in_axes)(*baked_data)
+            pause_memo()
+            try:
+                return jax.vmap(inner, in_axes=in_axes)(*baked_data)
+            finally:
+                resume_memo()
 
     return vmapped_fn

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev36
+Version: 0.5.3.dev37
 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.dev37/openscvx.egg-info/scm_version.json

@@ -0,0 +1,8 @@
+{
+  "tag": "0.5.2",
+  "distance": 37,
+  "node": "gdb86fef25495d5de3c07e7b227642a8e763af01d",
+  "dirty": false,
+  "branch": "main",
+  "node_date": "2026-07-09"
+}

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

@@ -1,56 +0,0 @@
-"""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
-
-from openscvx.symbolic.expr import Expr
-from openscvx.symbolic.lowerers.jax._registry import dispatch
-
-
-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.
-
-    Note:
-        This is a stateless lowerer - all state is in the expression tree.
-
-    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:
-        The lowerer is stateless and can be reused for multiple expressions.
-    """
-
-    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.
-
-        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)
-        """
-        return dispatch(self, expr)

openscvx-0.5.3.dev36/openscvx.egg-info/scm_version.json

@@ -1,8 +0,0 @@
-{
-  "tag": "0.5.2",
-  "distance": 36,
-  "node": "g1eddc1a14ee7e851448b965cc29df898aa2a822a",
-  "dirty": false,
-  "branch": "main",
-  "node_date": "2026-07-06"
-}