openscvx 0.5.3.dev35__tar.gz → 0.5.3.dev37__tar.gz

openscvx-0.5.3.dev37/.github/workflows/_extra.yml

@@ -0,0 +1,41 @@
+name: 🔌 Third-Party Extra
+
+# Reusable workflow behind the per-integration badge workflows (tests-mjx.yml,
+# tests-lie.yml, ...). Installs one optional extra and runs exactly the tests
+# carrying its marker — marker name == extra name, so one input drives both.
+
+on:
+  workflow_call:
+    inputs:
+      extra:
+        description: 'pyproject optional extra to install; also the pytest marker to select'
+        required: true
+        type: string
+
+jobs:
+  test:
+    name: ${{ inputs.extra }} (Python 3.13, ubuntu-latest)
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.13'
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v7
+      with:
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv pip install --system -e ".[test,${{ inputs.extra }}]"
+
+    - name: Run ${{ inputs.extra }} tests
+      run: |
+        pytest -v -m "${{ inputs.extra }} and not examples" --durations=10 --tb=short -n auto

openscvx-0.5.3.dev37/.github/workflows/tests-cvxpygen.yml

@@ -0,0 +1,18 @@
+name: cvxpygen
+
+# Badge workflow: health of the cvxpygen integration. Currently failing at
+# install (#551: cvxpygen>=1.0.0 needs cvxpy>=1.9, conflicting with the <1.9
+# pin) — cron + dispatch only while broken, so merges to main don't carry a
+# guaranteed red check; add a push trigger alongside the extras-job install
+# once the pins reconcile.
+
+on:
+  schedule:
+    - cron: '0 6 * * 1'  # Weekly on Monday, offset from the Sunday examples sweep
+  workflow_dispatch:
+
+jobs:
+  cvxpygen:
+    uses: ./.github/workflows/_extra.yml
+    with:
+      extra: cvxpygen

openscvx-0.5.3.dev35/.github/workflows/tests-integration.yml → openscvx-0.5.3.dev37/.github/workflows/tests-examples.yml

@@ -1,4 +1,4 @@
-name: 🧪 Integration Tests
+name: 🧪 Examples Sweep
 
 on:
   release:
@@ -8,10 +8,10 @@ on:
     - cron: '0 0 * * 0'  # Weekly on Sunday at midnight UTC
 
 jobs:
-  integration-tests:
-    name: Integration Tests
+  examples:
+    name: Examples Sweep
     runs-on: ${{ matrix.os }}
-    timeout-minutes: 30  # Integration tests can be slow
+    timeout-minutes: 30  # Solving every example is slow
 
     strategy:
       fail-fast: false  # Continue testing other examples even if one fails
@@ -35,18 +35,21 @@ jobs:
 
     - name: Install dependencies
       run: |
-        # Install package with optional cvxpygen and test packages
-        uv pip install --system -e ".[cvxpygen, test]"
+        # Install the package with every extra the examples exercise.
+        # cvxpygen omitted: its extra is unsatisfiable while cvxpy is pinned <1.9
+        # (cvxpygen>=1.0.0 needs cvxpy>=1.9, #551); its tests skip until that resolves.
+        # qpax omitted until its solves are stable on Linux runners (#550).
+        uv pip install --system -e ".[test,mjx,lie,frax,stl]"
 
-    - name: Run integration tests
+    - name: Run examples sweep
       run: |
-        pytest -v -m integration --durations=10 -n auto
+        pytest -v -m examples --durations=10 -n auto
 
     - name: Upload test results on failure
       if: failure()
       uses: actions/upload-artifact@v4
       with:
-        name: integration-test-results-py${{ matrix.python-version }}-${{ matrix.os }}
+        name: examples-test-results-py${{ matrix.python-version }}-${{ matrix.os }}
         path: |
           .pytest_cache/
           test-results/

openscvx-0.5.3.dev37/.github/workflows/tests-lie.yml

@@ -0,0 +1,18 @@
+name: jaxlie
+
+# Badge workflow: health of the jaxlie (Lie group) integration on main.
+# The weekly cron catches breakage from upstream releases between commits.
+
+on:
+  push:
+    branches:
+      - main
+  schedule:
+    - cron: '0 6 * * 1'  # Weekly on Monday, offset from the Sunday examples sweep
+  workflow_dispatch:
+
+jobs:
+  lie:
+    uses: ./.github/workflows/_extra.yml
+    with:
+      extra: lie

openscvx-0.5.3.dev37/.github/workflows/tests-mjx.yml

@@ -0,0 +1,18 @@
+name: mjx
+
+# Badge workflow: health of the MuJoCo MJX integration on main.
+# The weekly cron catches breakage from upstream releases between commits.
+
+on:
+  push:
+    branches:
+      - main
+  schedule:
+    - cron: '0 6 * * 1'  # Weekly on Monday, offset from the Sunday examples sweep
+  workflow_dispatch:
+
+jobs:
+  mjx:
+    uses: ./.github/workflows/_extra.yml
+    with:
+      extra: mjx

openscvx-0.5.3.dev37/.github/workflows/tests-qpax.yml

@@ -0,0 +1,16 @@
+name: qpax
+
+# Badge workflow: health of the qpax integration. Currently failing (#550) —
+# cron + dispatch only while broken, so merges to main don't carry a guaranteed
+# red check; add a push trigger alongside the extras-job install once fixed.
+
+on:
+  schedule:
+    - cron: '0 6 * * 1'  # Weekly on Monday, offset from the Sunday examples sweep
+  workflow_dispatch:
+
+jobs:
+  qpax:
+    uses: ./.github/workflows/_extra.yml
+    with:
+      extra: qpax

openscvx-0.5.3.dev37/.github/workflows/tests.yml

@@ -0,0 +1,140 @@
+name: 🧪 Tests
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+    branches:
+      - main
+    types: [opened, synchronize, reopened, ready_for_review]
+
+jobs:
+  unit:
+    if: ${{ !github.event.pull_request.draft }}
+    name: Unit Tests (Python ${{ matrix.python-version }}, ${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 15  # Unit tests should be fast
+
+    strategy:
+      fail-fast: false  # Continue testing other versions even if one fails
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ['3.11', '3.12', '3.13', '3.14']
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v7
+      with:
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv pip install --system -e ".[test]"
+
+    - name: Run unit tests
+      run: |
+        pytest -v -m "not e2e and not examples and not extras" --durations=10 --tb=short -n auto
+
+    - name: Upload test results on failure
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: unit-test-results-py${{ matrix.python-version }}-${{ matrix.os }}
+        path: |
+          .pytest_cache/
+          test-results/
+        retention-days: 7
+
+  e2e:
+    needs: unit
+    name: E2E Tests (Python ${{ matrix.python-version }}, ${{ matrix.os }})
+    runs-on: ${{ matrix.os }}
+    timeout-minutes: 15
+
+    strategy:
+      fail-fast: false  # Continue testing other versions even if one fails
+      matrix:
+        os: [ubuntu-latest, macos-latest]
+        python-version: ['3.11', '3.12', '3.13', '3.14']
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v7
+      with:
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        uv pip install --system -e ".[test]"
+
+    - name: Run e2e tests
+      run: |
+        pytest -v -m "e2e and not extras" --durations=10 --tb=short -n auto
+
+    - name: Upload test results on failure
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: e2e-test-results-py${{ matrix.python-version }}-${{ matrix.os }}
+        path: |
+          .pytest_cache/
+          test-results/
+        retention-days: 7
+
+  extras:
+    needs: unit
+    name: Extras Tests (Python 3.13, ubuntu-latest)
+    runs-on: ubuntu-latest
+    timeout-minutes: 15
+
+    steps:
+    - name: Checkout repository
+      uses: actions/checkout@v4
+
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: '3.13'
+
+    - name: Install uv
+      uses: astral-sh/setup-uv@v7
+      with:
+        enable-cache: true
+
+    - name: Install dependencies
+      run: |
+        # cvxpygen omitted: its extra is unsatisfiable while cvxpy is pinned <1.9
+        # (cvxpygen>=1.0.0 needs cvxpy>=1.9, #551); its tests skip until that resolves.
+        # qpax omitted until its solves are stable on Linux runners (#550).
+        uv pip install --system -e ".[test,mjx,lie]"
+
+    - name: Run extras tests
+      run: |
+        pytest -v -m "extras and not examples" --durations=10 --tb=short -n auto
+
+    - name: Upload test results on failure
+      if: failure()
+      uses: actions/upload-artifact@v4
+      with:
+        name: extras-test-results
+        path: |
+          .pytest_cache/
+          test-results/
+        retention-days: 7

{openscvx-0.5.3.dev35 → openscvx-0.5.3.dev37}/CONTRIBUTING.md

@@ -149,18 +149,19 @@ See the `[tool.ruff]` section in `pyproject.toml` for our configuration.
 
 OpenSCvx uses `pytest` for automated testing. Tests run automatically via GitHub Actions:
 
-- **[tests-unit.yml](.github/workflows/tests-unit.yml)**: Runs on every PR - fast unit tests including the brachistochrone benchmark
-- **[tests-integration.yml](.github/workflows/tests-integration.yml)**: Runs weekly - full example problems as a smoke test
+- **[tests.yml](.github/workflows/tests.yml)**: Runs on every PR - a fast `unit` job, then gated `e2e` (full SCP solves) and `extras` (optional-dependency tests) jobs
+- **[tests-examples.yml](.github/workflows/tests-examples.yml)**: Runs weekly - solves every example problem as a smoke test
 
 ##### Running Tests Locally
 
 For faster feedback during development, run tests locally:
 
 ```bash
-pytest tests/                              # All unit tests
-pytest tests/test_brachistochrone.py -v    # Just the brachistochrone benchmark
-pytest -v -m "not integration"             # Skip expensive integration tests
-pytest -v -m integration                   # Only integration tests (runs test_examples.py)
+pytest -m "not e2e and not examples and not extras"    # Fast unit tests
+pytest -m "e2e and not extras"                         # Full SCP-solve tests
+pytest -m extras                                       # Optional-dependency tests
+pytest -m examples                                     # Solve every example (slow)
+pytest tests/e2e/test_brachistochrone.py -v            # Just the brachistochrone benchmark
 ```
 
 ##### Brachistochrone Problem

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openscvx
-Version: 0.5.3.dev35
+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
@@ -41,11 +41,8 @@ Provides-Extra: frax
 Requires-Dist: frax>=0.0.4; extra == "frax"
 Provides-Extra: test
 Requires-Dist: pytest; extra == "test"
-Requires-Dist: scipy; extra == "test"
-Requires-Dist: jaxlie; extra == "test"
-Requires-Dist: svgpathtools; extra == "test"
 Requires-Dist: pytest-xdist; extra == "test"
-Requires-Dist: frax>=0.0.4; extra == "test"
+Requires-Dist: scipy; extra == "test"
 Dynamic: license-file
 
 <a id="readme-top"></a>
@@ -53,8 +50,8 @@ Dynamic: license-file
 <img src="figures/openscvx_logo.svg" width="1200"/>
 <p align="center">
     <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/lint.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/lint.yml/badge.svg"/></a>
-    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-unit.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-unit.yml/badge.svg"/></a>
-    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-integration.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-integration.yml/badge.svg"/></a>
+    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests.yml/badge.svg"/></a>
+    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-examples.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-examples.yml/badge.svg"/></a>
     <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/nightly.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/nightly.yml/badge.svg"/></a>
     <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/release.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/release.yml/badge.svg?event=release"/></a>
 </p>
@@ -253,6 +250,22 @@ This repo has the following features:
 
 <p align="right">(<a href="#readme-top">back to top</a>)</p>
 
+## Third-Party Integrations
+
+OpenSCvx integrates with several optional third-party packages. Each installs as a pip extra (`pip install openscvx[<extra>]`); without it, the library and its tests degrade gracefully — the corresponding tests simply skip. Integrations with a badge are tested against `main` on every push and re-checked weekly against the latest upstream releases.
+
+| Integration | Extra | Provides | Status |
+| --- | --- | --- | --- |
+| [MuJoCo MJX](https://mujoco.readthedocs.io/en/stable/mjx.html) | `mjx` | MJX dynamics adapter (`openscvx.integrations`) | [![mjx](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-mjx.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-mjx.yml) |
+| [jaxlie](https://github.com/brentyi/jaxlie) | `lie` | SO(3)/SE(3) Lie-group operations and IK initialization | [![jaxlie](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-lie.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-lie.yml) |
+| [qpax](https://github.com/qpax-solver/qpax) | `qpax` | JAX-native QP solver backend | [![qpax](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-qpax.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/issues/550) |
+| [CVXPYGen](https://github.com/cvxgrp/cvxpygen) | `cvxpygen` | Generated C solver code for the convex subproblem | [![cvxpygen](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-cvxpygen.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/issues/551) |
+| [moreau](https://pypi.org/project/moreau/) | `moreau` | Licensed QP solver backend | ![moreau](https://img.shields.io/badge/moreau-license%20required-lightgrey) |
+| [stljax](https://github.com/UW-CTRL/stljax) | `stl` | Signal Temporal Logic robustness bridge | ![stljax](https://img.shields.io/badge/stljax-no%20tests%20yet-lightgrey) |
+| [frax](https://github.com/danielpmorton/frax) | `frax` | Robot dynamics for the manipulator examples | [![frax](https://img.shields.io/badge/frax-weekly%20examples%20sweep-blue)](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-examples.yml) |
+
+<p align="right">(<a href="#readme-top">back to top</a>)</p>
+
 ## Acknowledgements
 
 This work was supported by a NASA Space Technology Graduate Research Opportunity and the Office of Naval Research under grant N00014-17-1-2433. The authors would like to acknowledge Natalia Pavlasek, Fabio Spada, Samuel Buckner, Abhi Kamath, Govind Chari, and Purnanand Elango as well as the other Autonomous Controls Laboratory members, for their many helpful discussions and support throughout this work.

{openscvx-0.5.3.dev35 → openscvx-0.5.3.dev37}/README.md

@@ -3,8 +3,8 @@
 <img src="figures/openscvx_logo.svg" width="1200"/>
 <p align="center">
     <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/lint.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/lint.yml/badge.svg"/></a>
-    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-unit.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-unit.yml/badge.svg"/></a>
-    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-integration.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-integration.yml/badge.svg"/></a>
+    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests.yml/badge.svg"/></a>
+    <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-examples.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-examples.yml/badge.svg"/></a>
     <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/nightly.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/nightly.yml/badge.svg"/></a>
     <a href="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/release.yml"><img src="https://github.com/OpenSCvx/OpenSCvx/actions/workflows/release.yml/badge.svg?event=release"/></a>
 </p>
@@ -203,6 +203,22 @@ This repo has the following features:
 
 <p align="right">(<a href="#readme-top">back to top</a>)</p>
 
+## Third-Party Integrations
+
+OpenSCvx integrates with several optional third-party packages. Each installs as a pip extra (`pip install openscvx[<extra>]`); without it, the library and its tests degrade gracefully — the corresponding tests simply skip. Integrations with a badge are tested against `main` on every push and re-checked weekly against the latest upstream releases.
+
+| Integration | Extra | Provides | Status |
+| --- | --- | --- | --- |
+| [MuJoCo MJX](https://mujoco.readthedocs.io/en/stable/mjx.html) | `mjx` | MJX dynamics adapter (`openscvx.integrations`) | [![mjx](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-mjx.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-mjx.yml) |
+| [jaxlie](https://github.com/brentyi/jaxlie) | `lie` | SO(3)/SE(3) Lie-group operations and IK initialization | [![jaxlie](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-lie.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-lie.yml) |
+| [qpax](https://github.com/qpax-solver/qpax) | `qpax` | JAX-native QP solver backend | [![qpax](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-qpax.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/issues/550) |
+| [CVXPYGen](https://github.com/cvxgrp/cvxpygen) | `cvxpygen` | Generated C solver code for the convex subproblem | [![cvxpygen](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-cvxpygen.yml/badge.svg?branch=main)](https://github.com/OpenSCvx/OpenSCvx/issues/551) |
+| [moreau](https://pypi.org/project/moreau/) | `moreau` | Licensed QP solver backend | ![moreau](https://img.shields.io/badge/moreau-license%20required-lightgrey) |
+| [stljax](https://github.com/UW-CTRL/stljax) | `stl` | Signal Temporal Logic robustness bridge | ![stljax](https://img.shields.io/badge/stljax-no%20tests%20yet-lightgrey) |
+| [frax](https://github.com/danielpmorton/frax) | `frax` | Robot dynamics for the manipulator examples | [![frax](https://img.shields.io/badge/frax-weekly%20examples%20sweep-blue)](https://github.com/OpenSCvx/OpenSCvx/actions/workflows/tests-examples.yml) |
+
+<p align="right">(<a href="#readme-top">back to top</a>)</p>
+
 ## Acknowledgements
 
 This work was supported by a NASA Space Technology Graduate Research Opportunity and the Office of Naval Research under grant N00014-17-1-2433. The authors would like to acknowledge Natalia Pavlasek, Fabio Spada, Samuel Buckner, Abhi Kamath, Govind Chari, and Purnanand Elango as well as the other Autonomous Controls Laboratory members, for their many helpful discussions and support throughout this work.

{openscvx-0.5.3.dev35 → openscvx-0.5.3.dev37}/docs/UsersGuide/01_hello_world_brachistochrone.md

@@ -58,7 +58,7 @@ $$
 t_f^* = \sqrt{\frac{r}{g}} \phi_f
 $$
 
-Because of these advantageous properties the Brachistochrone problem is extensively leveraged as a [unit test](https://github.com/OpenSCvx/OpenSCvx/blob/main/tests/test_brachistochrone.py).
+Because of these advantageous properties the Brachistochrone problem is extensively leveraged as a [unit test](https://github.com/OpenSCvx/OpenSCvx/blob/main/tests/e2e/test_brachistochrone.py).
 We would highly recommend that anyone setting out to develop some kind of optimization software do the same.
 It may not be necessary _nor_ sufficient, but a lot of things have to be working properly to solve Brachistochrone problem.
 

{openscvx-0.5.3.dev35 → 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.dev35'
-__version_tuple__ = version_tuple = (0, 5, 3, 'dev35')
+__version__ = version = '0.5.3.dev37'
+__version_tuple__ = version_tuple = (0, 5, 3, 'dev37')
 
-__commit_id__ = commit_id = 'g43f7b7f31'
+__commit_id__ = commit_id = 'gdb86fef25'

{openscvx-0.5.3.dev35 → 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