Skip to content

jax

JAX backend for lowering symbolic expressions to executable functions.

This package implements the JAX lowering backend that converts symbolic expression AST nodes into JAX functions with automatic differentiation support. The lowering uses a visitor pattern where each expression type has a corresponding visitor function registered via @visitor.

The visitor functions are split across submodules that mirror the openscvx.symbolic.expr package structure. Importing this package triggers registration of all visitors.

Example::

from openscvx.symbolic.lowerers.jax import JaxLowerer

lowerer = JaxLowerer()
f = lowerer.lower(expr)
result = f(x_val, u_val, node=0, params={})

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.

Source code in openscvx/symbolic/lowerers/jax/_lowerer.py 
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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
lower(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.

Parameters:

Name Type Description Default
expr Expr

Symbolic expression to lower (any Expr subclass)

 required

Returns:

Type Description
Callable

JAX function with signature (x, u, node, params) -> result

Raises:

Type Description
NotImplementedError

If no visitor exists for the expression type
ValueError

If the expression is malformed (e.g., State without slice)
Source code in openscvx/symbolic/lowerers/jax/_lowerer.py 
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
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