From 9114d386b94a20e87130bcda474dd93fc4947eb7 Mon Sep 17 00:00:00 2001
From: Michael Warner <hgxe14@gmail.com>
Date: Thu, 8 Oct 2026 15:58:14 -0700
Subject: [PATCH 1/4] Add csdl.experimental.compress for JAX-compiled subgraph
 operations

compress(inputs, outputs) moves the operations between the given variables
into one JaxCompressedOperation, a CompressedOperation subclass that:

- evaluates inline (recorder.inline, PySimulator) with one jax.jit function
  instead of executing the subgraph op by op;
- differentiates with CSDL's own reverse mode, recorded into a copy of the
  subgraph and wrapped as another JaxCompressedOperation, so the derivative
  is one compiled node, custom and implicit operations keep their CSDL
  derivatives, and higher derivatives work;
- by default (compile_separately=True) runs under JaxSimulator through a host
  callback as its own executable on its own device, which keeps it out of
  the simulator's XLA compile; with False it is traced into that program.

Options: absorb_feeders, intermediates ('none', 'used', 'all'), device, and
jit_kwargs. Backend-matrix tests are TestJaxCompress in
compress_operations.py; option, error, and exact-match tests are in
src/operations/tests/test_compress.py.
---
 csdl_alpha/experimental.py                    |   3 +
 .../src/operations/compress_operations.py     | 502 ++++++++++++++++++
 csdl_alpha/src/operations/tests/__init__.py   |   0
 .../src/operations/tests/test_compress.py     | 502 ++++++++++++++++++
 4 files changed, 1007 insertions(+)
 create mode 100644 csdl_alpha/src/operations/tests/__init__.py
 create mode 100644 csdl_alpha/src/operations/tests/test_compress.py

diff --git a/csdl_alpha/experimental.py b/csdl_alpha/experimental.py
index efa57ca..73035a2 100644
--- a/csdl_alpha/experimental.py
+++ b/csdl_alpha/experimental.py
@@ -5,6 +5,9 @@
 # Tensor
 from csdl_alpha.src.operations.tensor.batch import batch_function
 
+# Compress a subgraph into one JAX-compiled operation
+from csdl_alpha.src.operations.compress_operations import compress, JaxCompressedOperation
+
 # Simulator API
 from csdl_alpha.backends.simulator import PySimulator
 from csdl_alpha.backends.jax.jax_simulator import JaxSimulator
diff --git a/csdl_alpha/src/operations/compress_operations.py b/csdl_alpha/src/operations/compress_operations.py
index 3db1628..8884c60 100644
--- a/csdl_alpha/src/operations/compress_operations.py
+++ b/csdl_alpha/src/operations/compress_operations.py
@@ -109,6 +109,379 @@ def compress_current_operations():
     compressed_operation.finalize_and_return_outputs(skip_inline=True)
     return compressed_operation
 
+def _resolve_device(device):
+    """``None``, a ``jax.Device``, or a platform name (``'cpu'``, ``'gpu'``) -> ``jax.Device`` or ``None``."""
+    if device is None or not isinstance(device, str):
+        return device
+    import jax
+    try:
+        return jax.devices(device)[0]
+    except RuntimeError as error:
+        raise ValueError(f"No JAX device for {device!r}: {error}") from None
+
+
+class JaxCompressedOperation(CompressedOperation):
+    """A ``CompressedOperation`` that runs as one jitted JAX function, derivatives included.
+
+    The subgraph is translated with ``create_jax_function`` and wrapped in
+    ``jax.jit(..., **jit_kwargs)`` once, on first use. Compared with
+    ``CompressedOperation``:
+
+    - Inline evaluation (``recorder.inline``, ``PySimulator``) calls that
+      jitted function on ``device`` instead of executing the subgraph
+      operation by operation.
+    - Inside an enclosing JAX trace (``JaxSimulator``, an enclosing compressed
+      operation), ``compute_jax`` either calls the function through a host
+      callback, so that it stays its own executable on its own ``device``
+      (``compile_separately=True``, the default), or traces it into the
+      enclosing program, which XLA then compiles and optimizes as a whole.
+    - Reverse mode records CSDL's own VJP of the subgraph into a copy of it
+      and wraps that copy as another ``JaxCompressedOperation``. Every
+      operation inside uses its CSDL derivative (custom operations call
+      ``compute_derivatives``; implicit operations use their nonlinear
+      solver's implicit derivative), the derivative is one compiled node on
+      the graph, and it is differentiable in turn.
+
+    Parameters
+    ----------
+    subgraph : Graph
+        The graph to run. Its leaf variables are ``inputs`` or literal
+        constants, which are compiled in.
+    inputs, outputs : list[Variable]
+        Variables of ``subgraph`` that become this operation's inputs and
+        outputs. Repeated inputs are merged.
+    name : str
+        Operation name.
+    device : str or jax.Device, optional
+        Where the operation runs (``'cpu'``, ``'gpu'``, or a ``jax.Device``);
+        by default JAX's default device. Inside an enclosing JAX trace this
+        applies only with ``compile_separately``.
+    jit_kwargs : dict, optional
+        Extra keyword arguments for ``jax.jit``.
+    compile_separately : bool
+        Keep this operation out of any enclosing JAX program, by default True.
+        The enclosing program then compiles without it, which cuts its compile
+        time and memory, and this operation runs on ``device``. Each call costs
+        a round trip through host memory and loses XLA optimization across the
+        boundary; under ``vmap`` the callback runs once per batch element. With
+        False, the operation is traced into the enclosing program instead.
+
+    Derivative operations inherit ``device``, ``jit_kwargs``, and
+    ``compile_separately``.
+    """
+
+    def __init__(
+            self,
+            subgraph,
+            inputs,
+            outputs,
+            name = 'jax_compressed',
+            device = None,
+            jit_kwargs = None,
+            compile_separately = True,
+        ):
+        import jax
+        jax.config.update("jax_enable_x64", True)
+        super().__init__(subgraph, list(dict.fromkeys(inputs)), list(outputs), name)
+        self.device = _resolve_device(device)
+        self.jit_kwargs = dict(jit_kwargs or {})
+        self.compile_separately = compile_separately
+
+    @property
+    def jit_fn(self):
+        """The subgraph as a ``jax.jit`` function, built on first use."""
+        if self.jax_function is None:
+            import jax
+            from csdl_alpha.backends.jax.graph_to_jax import create_jax_function
+            fn = create_jax_function(self.get_subgraph(), self.outputs, self.inputs)
+            self.jax_function = jax.jit(fn, **self.jit_kwargs)
+        return self.jax_function
+
+    def compute_jax(self, *args):
+        if not self.compile_separately:
+            return tuple(self.jit_fn(*args))
+        import jax
+        from csdl_alpha.backends.jax.utils import host_callback
+        result_shapes = [jax.ShapeDtypeStruct(var.shape, np.float64) for var in self.outputs]
+        return tuple(host_callback(self._compute_on_host, result_shapes, *args))
+
+    def _compute_on_host(self, *args):
+        outs = self.compute_inline(*args)
+        return (outs,) if self.num_outputs == 1 else outs
+
+    def compute_inline(self, *args):
+        if self.device is not None:
+            import jax
+            args = jax.device_put(args, self.device)
+        outs = self.jit_fn(*args)
+        outs = tuple(np.array(out).reshape(var.shape) for out, var in zip(outs, self.outputs))
+        if self.num_outputs == 1:
+            return outs[0]
+        return outs
+
+    def evaluate_vjp(self, cotangents, *inputs_outputs):
+        import csdl_alpha as csdl
+        from csdl_alpha.src.operations.derivatives.reverse import vjp
+
+        inputs = inputs_outputs[:self.num_inputs]
+        outputs = inputs_outputs[self.num_inputs:]
+        # Outputs with a None cotangent contribute nothing; leave them out.
+        seeds = [(y, cotangents[y]) for y in outputs if cotangents.check(y) and cotangents[y] is not None]
+        wrts = [x for x in inputs if cotangents.check(x)]
+        if not seeds or not wrts:
+            return
+
+        # Record CSDL's reverse mode into a copy of the subgraph, so the subgraph
+        # itself stays a clean forward graph for later derivative calls. The
+        # recorder creates the copy so that it sits in its graph tree, which
+        # operations that open their own subgraphs (loops, composed ops) need.
+        # The outer cotangent variables join the copy as its seed inputs.
+        recorder = csdl.get_current_recorder()
+        inline = recorder.inline
+        recorder.inline = False
+        recorder._enter_subgraph(name=f'vjp_{self.name}')
+        work = recorder.active_graph
+        work.rxgraph = self.get_subgraph().rxgraph.copy()
+        work.update_node_table()
+        try:
+            for _, seed in seeds:
+                work.add_node(seed)
+            input_cotangents = vjp(seeds, wrts, work)
+            # Fresh variables, so each gradient is produced inside the copy.
+            grads = [csdl.copyvar(input_cotangents[x]) if input_cotangents[x] is not None
+                     else csdl.copyvar(csdl.Variable(value=np.zeros(x.shape))) for x in wrts]
+        finally:
+            recorder._exit_subgraph()
+            recorder.inline = inline
+
+        for grad in grads:
+            recorder._add_node(grad)
+        # The copy also holds forward operations the VJP does not need; jax.jit drops them.
+        vjp_op = JaxCompressedOperation(
+            work, list(inputs) + [seed for _, seed in seeds], grads, name=f'vjp_{self.name}',
+            device=self.device, jit_kwargs=self.jit_kwargs, compile_separately=self.compile_separately)
+        vjp_op.finalize_and_return_outputs()
+        for x, grad in zip(wrts, grads):
+            cotangents.accumulate(x, grad)
+
+
+def compress(
+        inputs,
+        outputs,
+        name = 'jax_compressed',
+        absorb_feeders = False,
+        intermediates = 'none',
+        device = None,
+        jit_kwargs = None,
+        compile_separately = True,
+    ):
+    """Replace the operations between ``inputs`` and ``outputs`` with one JAX-compiled operation.
+
+    The compressed region is every operation that both depends on ``inputs``
+    and feeds ``outputs`` (plus, with ``absorb_feeders``, upstream operations
+    that only feed it). Those operations and their intermediate variables move
+    out of the active graph into the subgraph of a single
+    :class:`JaxCompressedOperation`. That operation runs the region as one
+    ``jax.jit`` function and differentiates it with CSDL's own reverse mode,
+    compiled the same way. The ``outputs`` variables are kept: the new operation
+    produces them, so downstream operations, constraints, objectives, and
+    simulator handles still work.
+
+    Any other variable the region reads (a parameter, or a value computed
+    upstream) becomes an extra input whose value is read from the graph, so it
+    is never frozen and is differentiable. Only literal constants
+    (``csdl.Constant``, such as the ``2.0`` in ``2.0 * x``) are compiled in.
+
+    Parameters
+    ----------
+    inputs : Variable or list[Variable]
+        Where the region starts. No input may depend on another.
+    outputs : Variable or list[Variable]
+        Where the region ends.
+    name : str
+        Name of the new operation.
+    absorb_feeders : bool
+        Also compress upstream operations whose outputs only the region uses,
+        such as the broadcast of a scalar constant, by default False. Each one
+        left outside stays a separate operation.
+    intermediates : str
+        Which intermediates also become outputs. ``'none'`` (default): only
+        ``outputs``; an intermediate used outside the region (or registered as
+        a constraint or objective) raises. ``'used'``: add those
+        intermediates as outputs. ``'all'``: every variable the region
+        computes is an output.
+    device : str or jax.Device, optional
+        Where the operation runs (``'cpu'``, ``'gpu'``, or a ``jax.Device``).
+        Under ``JaxSimulator`` this applies only with ``compile_separately``.
+    jit_kwargs : dict, optional
+        Extra keyword arguments for ``jax.jit``.
+    compile_separately : bool
+        Under ``JaxSimulator``, call this operation's own compiled function
+        through a host callback instead of compiling it into the simulator's
+        program, by default True. Set False to let XLA compile and optimize
+        the region together with the rest of the program. See
+        :class:`JaxCompressedOperation`.
+
+    Returns
+    -------
+    JaxCompressedOperation
+        ``op.inputs`` is ``inputs`` followed by ``op.extra_inputs``,
+        ``op.outputs`` is ``outputs`` followed by ``op.extra_outputs``
+        (intermediates added by ``intermediates``), and ``op.get_subgraph()``
+        is the region.
+
+    Examples
+    --------
+    >>> recorder = csdl.Recorder(inline=True)
+    >>> recorder.start()
+    >>> x = csdl.Variable(value=np.array([0.1, 0.2]))
+    >>> z = csdl.sum(csdl.sin(x) * x)
+    >>> op = csdl.experimental.compress(x, z)
+    >>> f = 2.0 * z
+    >>> csdl.derivative(f, x).value
+    array([[0.39866767, 0.78936529]])
+    """
+    import csdl_alpha as csdl
+    from csdl_alpha.src.graph.variable import Constant
+    from csdl_alpha.utils.inputs import listify_variables
+
+    if intermediates not in ('none', 'used', 'all'):
+        raise ValueError(f"intermediates must be 'none', 'used', or 'all', not {intermediates!r}.")
+    device = _resolve_device(device)
+    inputs = listify_variables(inputs)
+    outputs = listify_variables(outputs)
+    for variables, label in [(inputs, 'inputs'), (outputs, 'outputs')]:
+        if len(set(variables)) != len(variables):
+            raise ValueError(f"`{label}` contains a variable more than once.")
+
+    recorder = csdl.get_current_recorder()
+    graph = recorder.active_graph
+    for var in inputs + outputs:
+        if var not in graph.node_table:
+            raise ValueError(f"{var.info()} is not in the active graph.")
+    for var in outputs:
+        if var in inputs:
+            raise ValueError(f"{var.info()} is given as both an input and an output.")
+
+    ops, internal, extra_outputs = _find_region(graph, inputs, outputs, absorb_feeders, intermediates)
+    outputs = outputs + extra_outputs
+
+    # Gather the variables the region reads that it does not compute. Literal
+    # constants cannot change, so they are compiled into the function;
+    # everything else becomes an input.
+    input_set = set(inputs)
+    extra_inputs = []
+    constants = []
+    seen = set()
+    for op in ops:
+        for var in op.inputs:
+            if var not in internal and var not in input_set and var not in seen:
+                seen.add(var)
+                is_literal = type(var) is Constant and var.value is not None
+                (constants if is_literal else extra_inputs).append(var)
+    op_inputs = inputs + extra_inputs
+
+    # Move the region's operations into their own graph (this deletes them from
+    # `graph` and keeps the variables there), then drop the intermediates and
+    # any constants nothing else uses. The outputs stay; the new operation
+    # produces them.
+    region = graph.extract_subgraph_nodes(set(ops) | internal | set(op_inputs) | set(constants))
+    region.name = name
+    output_set = set(outputs)
+    graph._delete_nodes(
+        [var for var in internal if var not in output_set]
+        + [var for var in constants if graph.out_degree(var) == 0]
+    )
+
+    compressed = JaxCompressedOperation(
+        region, op_inputs, outputs, name=name, device=device,
+        jit_kwargs=jit_kwargs, compile_separately=compile_separately)
+    compressed.extra_inputs = extra_inputs
+    compressed.extra_outputs = extra_outputs
+    # Skip inline evaluation: the outputs already hold the values the region computed.
+    compressed.finalize_and_return_outputs(skip_inline=True)
+    return compressed
+
+
+def _find_region(graph, inputs, outputs, absorb_feeders, intermediates):
+    """Return the region's operations in topological order, the variables they
+    compute, and the intermediates to add as outputs."""
+    import csdl_alpha as csdl
+    import rustworkx as rx
+    from csdl_alpha.src.operations.operation_subclasses import RandomOperation
+
+    rxgraph = graph.rxgraph
+    input_indices = {graph.node_table[var] for var in inputs}
+    for var in inputs:
+        if input_indices & rx.ancestors(rxgraph, graph.node_table[var]):
+            raise ValueError(f"Input {var.info()} depends on another input.")
+
+    # Descendants of the inputs that are ancestors of the outputs. This raises
+    # if an input affects no output or an output depends on no input.
+    between = graph._get_intersection(
+        inputs, outputs, add_hanging_input_variables=False, add_hanging_output_variables=False)
+    op_indices = {i for i in between if isinstance(rxgraph[i], Operation)}
+
+    recorder = csdl.get_current_recorder()
+    registered = set(recorder.constraints) | set(recorder.objectives)
+    if absorb_feeders:
+        _absorb_feeders(graph, op_indices, input_indices, registered)
+
+    order = rx.topological_sort(rxgraph)
+    ops = [rxgraph[i] for i in order if i in op_indices]
+    internal = {var for op in ops for var in op.outputs}
+
+    for op in ops:
+        if isinstance(op, RandomOperation):
+            raise NotImplementedError(f"Cannot compress random operation {op.info()}.")
+
+    # An intermediate used outside the region would lose its producer, so it
+    # must become an output.
+    output_set = set(outputs)
+    extra_outputs = []
+    for var in (var for op in ops for var in op.outputs):
+        if var in output_set or var in extra_outputs:
+            continue
+        if intermediates == 'all':
+            extra_outputs.append(var)
+            continue
+        outside = [rxgraph[i] for i in rxgraph.successor_indices(graph.node_table[var]) if i not in op_indices]
+        if outside or var in registered:
+            if intermediates == 'used':
+                extra_outputs.append(var)
+                continue
+            users = ', '.join(op.info() for op in outside) or 'a constraint or objective'
+            raise ValueError(
+                f"Intermediate {var.info()} is used outside the compressed region (by {users}). "
+                "Add it to `outputs`, or pass intermediates='used'."
+            )
+    return ops, internal, extra_outputs
+
+
+def _absorb_feeders(graph, op_indices, input_indices, registered):
+    """Add to ``op_indices`` every upstream operation whose outputs are used only inside the region."""
+    rxgraph = graph.rxgraph
+
+    def producers_of_inputs(op_index):
+        for var_index in rxgraph.predecessor_indices(op_index):
+            if var_index not in input_indices:
+                yield from rxgraph.predecessor_indices(var_index)
+
+    stack = [p for i in op_indices for p in producers_of_inputs(i)]
+    while stack:
+        candidate = stack.pop()
+        if candidate in op_indices:
+            continue
+        exclusive = all(
+            rxgraph[out] not in registered
+            and all(consumer in op_indices for consumer in rxgraph.successor_indices(out))
+            for out in rxgraph.successor_indices(candidate)
+        )
+        if exclusive:
+            op_indices.add(candidate)
+            stack.extend(producers_of_inputs(candidate))
+
+
 class TestCompressOp(csdl_tests.CSDLTest):
     @staticmethod
     def simple_model(x, y):
@@ -134,6 +507,135 @@ def test_simple(self):
         compare_values += [csdl_tests.TestingPair(z, z_np, tag = 'simple')]
         self.run_tests(compare_values, verify_derivatives=True)
 
+
+class TestJaxCompress(csdl_tests.CSDLTest):
+    """compress() under every backend and derivative mode of the test matrix."""
+
+    x_val = np.array([0.1, 0.5, -0.3, 1.2])
+    y_val = np.array([1.0, 2.0, 0.5, -1.0])
+    c_val = np.array([0.7, -0.2])
+
+    @staticmethod
+    def region(x, y, c, lib):
+        """Elementwise ops, a matvec and a reduction; `lib` is csdl or a numpy shim."""
+        a = lib.sin(x) * y + lib.exp(0.3 * x)
+        b = lib.matvec(lib.reshape(a, (2, 2)), c)
+        return a, lib.sum(b**2) + lib.norm(a)
+
+    class np_lib:
+        sin, exp, sum = np.sin, np.exp, np.sum
+        reshape = staticmethod(np.reshape)
+        matvec = staticmethod(lambda A, v: A @ v)
+        norm = staticmethod(lambda a: np.array([np.linalg.norm(a)]))
+
+    def build(self, **compress_kwargs):
+        import csdl_alpha as csdl
+        x = csdl.Variable(name='x', value=self.x_val)
+        y = csdl.Variable(name='y', value=self.y_val)
+        c = csdl.Variable(name='c', value=self.c_val)  # read by the region but not passed to compress
+        a, z = self.region(x, y, c, csdl)
+        return x, y, c, a, z
+
+    def expected(self):
+        a, z = self.region(self.x_val, self.y_val, self.c_val, self.np_lib)
+        return a, np.array([2.0 * z.item() + np.sum(self.x_val)])
+
+    def test_docstring(self):
+        self.docstest(compress)
+
+    def test_options(self):
+        option_sets = [
+            {},
+            {'absorb_feeders': True},
+            {'compile_separately': False},
+            {'absorb_feeders': True, 'compile_separately': False, 'jit_kwargs': {'keep_unused': True}},
+        ]
+        import csdl_alpha as csdl
+        for options in option_sets:
+            self.prep()
+            x, y, c, a, z = self.build()
+            op = compress([x, y], z, **options)
+            assert isinstance(op, JaxCompressedOperation) and op.outputs == [z]
+            f = 2.0 * z + csdl.sum(x)
+            _, f_np = self.expected()
+            self.run_tests([csdl_tests.TestingPair(f, f_np, tag=str(options))], verify_derivatives=True)
+
+    def test_intermediates_used_outside(self):
+        import csdl_alpha as csdl
+        self.prep()
+        x, y, c, a, z = self.build()
+        w = csdl.sum(a * 3.0)  # uses the intermediate `a` outside x, y -> z
+        op = compress([x, y], z, intermediates='used')
+        assert op.extra_outputs == [a]
+        a_np, f_np = self.expected()
+        self.run_tests([
+            csdl_tests.TestingPair(2.0 * z + csdl.sum(x), f_np, tag='f'),
+            csdl_tests.TestingPair(w, np.array([3.0 * np.sum(a_np)]), tag='w'),
+        ], verify_derivatives=True)
+
+    def test_implicit_operation(self):
+        import csdl_alpha as csdl
+        a_val = np.array([2.0, 3.0])
+        for compile_separately in [False, True]:
+            self.prep()
+            a = csdl.Variable(name='a', value=a_val)
+            x = csdl.ImplicitVariable(name='x', value=np.ones(2))
+            solver = csdl.nonlinear_solvers.Newton(print_status=False, tolerance=1e-12)
+            solver.add_state(x, x**3 + a * x - 10.0)
+            solver.run()
+            z = csdl.sum(x**2 * a)
+            compress(a, z, compile_separately=compile_separately)
+            x_np = np.array([np.real(r[np.isreal(r)][0]) for r in (np.roots([1, 0, ai, -10]) for ai in a_val)])
+            self.run_tests(
+                [csdl_tests.TestingPair(z, np.array([np.sum(x_np**2 * a_val)]), decimal=9)],
+                verify_derivatives=True)
+
+    def test_custom_operation(self):
+        import csdl_alpha as csdl
+
+        class CubePlusProduct(csdl.CustomExplicitOperation):
+            def evaluate(self, x, y):
+                self.declare_input('x', x)
+                self.declare_input('y', y)
+                f = self.create_output('f', x.shape)
+                self.declare_derivative_parameters('f', 'x')
+                self.declare_derivative_parameters('f', 'y')
+                return f
+
+            def compute(self, input_vals, output_vals):
+                output_vals['f'] = input_vals['x']**3 + input_vals['x'] * input_vals['y']
+
+            def compute_derivatives(self, input_vals, output_vals, derivatives):
+                derivatives['f', 'x'] = np.diag(3 * input_vals['x']**2 + input_vals['y'])
+                derivatives['f', 'y'] = np.diag(input_vals['x'])
+
+        x_val, y_val = np.array([0.3, -0.4, 0.9]), np.array([1.0, 0.5, -2.0])
+        self.prep()
+        x = csdl.Variable(name='x', value=x_val)
+        y = csdl.Variable(name='y', value=y_val)
+        u = csdl.sin(x) * 2.0
+        z = csdl.sum(csdl.exp(0.1 * CubePlusProduct().evaluate(u, y)) * u)
+        compress([x, y], z)
+        u_np = np.sin(x_val) * 2.0
+        z_np = np.sum(np.exp(0.1 * (u_np**3 + u_np * y_val)) * u_np)
+        self.run_tests([csdl_tests.TestingPair(z, np.array([z_np]))], verify_derivatives=True)
+
+    def test_loop_and_nested(self):
+        import csdl_alpha as csdl
+        x_val = np.array([0.2, 0.4, 0.8])
+        self.prep()
+        x = csdl.Variable(name='x', value=x_val)
+        acc = csdl.Variable(value=np.zeros(3))
+        for i in csdl.frange(3):
+            acc = acc + csdl.sin(x * (i + 1.0))
+        z = csdl.sum(acc * x)
+        compress(x, z)                                         # inner, separate
+        f = csdl.exp(z * 0.1)
+        compress(x, f, name='outer', compile_separately=False)  # fused, contains the inner
+        z_np = np.sum(sum(np.sin(x_val * (i + 1.0)) for i in range(3)) * x_val)
+        self.run_tests([csdl_tests.TestingPair(f, np.array([np.exp(0.1 * z_np)]))], verify_derivatives=True)
+
+
 if __name__ == '__main__':
     test = TestCompressOp()
     test.overwrite_backend = 'jax'
diff --git a/csdl_alpha/src/operations/tests/__init__.py b/csdl_alpha/src/operations/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/csdl_alpha/src/operations/tests/test_compress.py b/csdl_alpha/src/operations/tests/test_compress.py
new file mode 100644
index 0000000..a118c49
--- /dev/null
+++ b/csdl_alpha/src/operations/tests/test_compress.py
@@ -0,0 +1,502 @@
+"""Option, error, and exact-match tests for csdl.experimental.compress.
+
+Each test compares a compressed model with the same model left uncompressed.
+The backend-matrix tests (values and finite-difference derivatives) are
+TestJaxCompress in compress_operations.py.
+"""
+import numpy as np
+import pytest
+
+import csdl_alpha as csdl
+from csdl_alpha.src.operations.operation_subclasses import SubgraphOperation
+from csdl_alpha.experimental import JaxCompressedOperation, compress
+
+
+def model(x, y, c):
+    """Region to compress: a few elementwise ops, a matvec, and a reduction."""
+    a = csdl.sin(x) * y + csdl.exp(0.3 * x)
+    A = csdl.reshape(a, (2, 2))
+    b = csdl.matvec(A, c)
+    z = csdl.sum(b**2) + csdl.norm(a)
+    return a, z
+
+
+def build(do_compress, inline=True, compress_outputs=('z',), **compress_kwargs):
+    rec = csdl.Recorder(inline=inline)
+    rec.start()
+    x = csdl.Variable(name='x', value=np.array([0.1, 0.5, -0.3, 1.2]))
+    y = csdl.Variable(name='y', value=np.array([1.0, 2.0, 0.5, -1.0]))
+    c = csdl.Variable(name='c', value=np.array([0.7, -0.2]))  # not passed to compress
+    a, z = model(x, y, c)
+    named = {'a': a, 'z': z}
+    op = compress([x, y], [named[k] for k in compress_outputs], **compress_kwargs) if do_compress else None
+    # downstream of the region, outside it
+    f = z * 2.0 + csdl.sum(x)
+    return rec, dict(x=x, y=y, c=c, a=a, z=z, f=f, op=op)
+
+
+def op_names(rec):
+    from csdl_alpha.src.graph.operation import Operation
+    return sorted(n.name for n in rec.active_graph.node_table if isinstance(n, Operation))
+
+
+def test_feeders_stay_outside_by_default():
+    ref_rec, ref = build(False)
+    ref_d = csdl.derivative(ref['f'], [ref['x'], ref['c']])
+    ref_rec.stop()
+
+    rec, v = build(True)
+    d = csdl.derivative(v['f'], [v['x'], v['c']])
+    rec.stop()
+    names = [n for n in op_names(rec) if not n.startswith('vjp')]
+    assert names.count('reshape') == 2
+    assert v['c'] not in v['op'].inputs
+    for k in ['x', 'c']:
+        np.testing.assert_allclose(d[v[k]].value, ref_d[ref[k]].value, rtol=1e-12)
+
+
+@pytest.mark.parametrize('loop', [True, False])
+def test_values_and_first_derivatives(loop):
+    ref_rec, ref = build(False)
+    ref_d = csdl.derivative(ref['f'], [ref['x'], ref['y'], ref['c']], loop=loop)
+    ref_rec.stop()
+
+    rec, v = build(True)
+    d = csdl.derivative(v['f'], [v['x'], v['y'], v['c']], loop=loop)
+    rec.stop()
+
+    np.testing.assert_allclose(v['f'].value, ref['f'].value)
+    for k in ['x', 'y', 'c']:
+        np.testing.assert_allclose(d[v[k]].value, ref_d[ref[k]].value, rtol=1e-12)
+
+
+def test_absorbed_feeders_make_region_one_node():
+    rec, v = build(True, absorb_feeders=True)
+    rec.stop()
+    op = v['op']
+    assert isinstance(op, JaxCompressedOperation)
+    # The reshapes of `c` and of the literal 0.3 feed only the region, so they are
+    # absorbed; the literal is compiled in. Left outside: the downstream ops.
+    assert op_names(rec) == ['add', 'jax_compressed', 'mult', 'sum']
+    assert op.inputs == [v['x'], v['y'], v['c']]
+    assert op.extra_inputs == [v['c']]
+    assert op.outputs == [v['z']]
+    assert v['a'] not in rec.active_graph.node_table
+
+
+def test_second_derivatives():
+    ref_rec, ref = build(False)
+    g = csdl.derivative(ref['f'], ref['x'])
+    H_ref = csdl.derivative(g, [ref['x'], ref['c']])
+    ref_rec.stop()
+
+    rec, v = build(True)
+    g = csdl.derivative(v['f'], v['x'])
+    H = csdl.derivative(g, [v['x'], v['c']])
+    rec.stop()
+
+    np.testing.assert_allclose(H[v['x']].value, H_ref[ref['x']].value, rtol=1e-10, atol=1e-12)
+    np.testing.assert_allclose(H[v['c']].value, H_ref[ref['c']].value, rtol=1e-10, atol=1e-12)
+
+
+def test_multiple_outputs_including_reused_intermediate():
+    # `a` is both an output and an input to the computation of `z`.
+    ref_rec, ref = build(False)
+    ref_g = csdl.derivative(ref['f'] + csdl.sum(ref['a']**3), ref['x'])
+    ref_rec.stop()
+
+    rec, v = build(True, compress_outputs=('a', 'z'))
+    g = csdl.derivative(v['f'] + csdl.sum(v['a']**3), v['x'])
+    rec.stop()
+
+    assert v['op'].outputs == [v['a'], v['z']]
+    np.testing.assert_allclose(g.value, ref_g.value, rtol=1e-12)
+
+
+def test_pysimulator_matches_finite_difference():
+    rec, v = build(True, inline=False)
+    rec.stop()
+    sim = csdl.experimental.PySimulator(rec)
+    sim[v['x']] = np.array([0.3, -0.1, 0.9, 0.2])
+    sim.run()
+
+    ref_rec, ref = build(False, inline=False)
+    ref_rec.stop()
+    ref_sim = csdl.experimental.PySimulator(ref_rec)
+    ref_sim[ref['x']] = np.array([0.3, -0.1, 0.9, 0.2])
+    ref_sim.run()
+    np.testing.assert_allclose(sim[v['f']], ref_sim[ref['f']])
+
+    wrts = [v['x'], v['y'], v['c']]
+    exact = sim.compute_totals(v['f'], wrts)
+    fd = sim.compute_totals(v['f'], wrts, use_finite_difference=True)
+    for w in wrts:
+        np.testing.assert_allclose(exact[v['f'], w], fd[v['f'], w], rtol=1e-5, atol=1e-6)
+
+
+def test_jax_simulator():
+    rec, v = build(True, inline=False)
+    rec.stop()
+    sim = csdl.experimental.JaxSimulator(
+        rec, gpu=False, additional_inputs=[v['x'], v['c']], additional_outputs=[v['f']])
+    sim[v['x']] = np.array([0.3, -0.1, 0.9, 0.2])
+    sim.run()
+
+    ref_rec, ref = build(False, inline=False)
+    ref_rec.stop()
+    ref_sim = csdl.experimental.JaxSimulator(
+        ref_rec, gpu=False, additional_inputs=[ref['x'], ref['c']], additional_outputs=[ref['f']])
+    ref_sim[ref['x']] = np.array([0.3, -0.1, 0.9, 0.2])
+    ref_sim.run()
+    np.testing.assert_allclose(sim[v['f']], ref_sim[ref['f']])
+
+    d = sim.compute_totals()
+    d_ref = ref_sim.compute_totals()
+    for k in ['x', 'c']:
+        np.testing.assert_allclose(d[v['f'], v[k]], d_ref[ref['f'], ref[k]], rtol=1e-12)
+
+
+def test_nested_compress():
+    ref_rec, ref = build(False)
+    ref_d = csdl.derivative(ref['f'], [ref['x'], ref['c']])
+    ref_rec.stop()
+
+    rec, v = build(True)       # inner compress: x, y -> z
+    outer = compress([v['x'], v['c']], v['f'], name='outer')
+    d = csdl.derivative(v['f'], [v['x'], v['c']])
+    rec.stop()
+
+    assert 'jax_compressed' not in op_names(rec)  # the inner op was absorbed into `outer`
+    assert outer.inputs[:2] == [v['x'], v['c']] and v['y'] in outer.inputs
+    np.testing.assert_allclose(d[v['x']].value, ref_d[ref['x']].value, rtol=1e-12)
+    np.testing.assert_allclose(d[v['c']].value, ref_d[ref['c']].value, rtol=1e-12)
+
+
+def test_parameter_change_flows_through_extra_input():
+    rec, v = build(True, inline=False)
+    rec.stop()
+    sim = csdl.experimental.PySimulator(rec)
+    sim.run()
+    f0 = sim[v['f']].copy()
+    sim[v['c']] = np.array([3.0, 1.0])  # c was never passed to compress
+    sim.run()
+    assert not np.allclose(sim[v['f']], f0)
+
+
+def test_loop_inside_region():
+    def body(x, use_loop):
+        if use_loop:
+            acc = csdl.Variable(value=np.zeros(x.shape))
+            for i in csdl.frange(3):
+                acc = acc + csdl.sin(x * (i + 1.0))
+        else:
+            acc = sum(csdl.sin(x * float(i + 1)) for i in range(3))
+        return csdl.sum(acc * x)
+
+    results = []
+    for do_compress in [False, True]:
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        x = csdl.Variable(name='x', value=np.array([0.2, 0.4, 0.8]))
+        z = body(x, use_loop=True)
+        if do_compress:
+            compress(x, z)
+        g = csdl.derivative(z, x)
+        rec.stop()
+        results.append((z.value, g.value))
+    np.testing.assert_allclose(results[1][0], results[0][0])
+    np.testing.assert_allclose(results[1][1], results[0][1], rtol=1e-12)
+
+
+def test_escaping_intermediate_raises():
+    rec = csdl.Recorder(inline=True)
+    rec.start()
+    x = csdl.Variable(value=np.ones(3))
+    a = csdl.sin(x)
+    z = csdl.sum(a)
+    w = a * 2.0  # uses the intermediate `a` outside x -> z
+    with pytest.raises(ValueError, match="intermediates='used'"):
+        compress(x, z)
+    rec.stop()
+
+
+@pytest.mark.parametrize('intermediates', ['used', 'all'])
+def test_intermediates_become_outputs(intermediates):
+    results = []
+    for do_compress in [False, True]:
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        x = csdl.Variable(value=np.array([0.2, 0.7, 1.1]))
+        a = csdl.sin(x)
+        b = csdl.exp(a)
+        z = csdl.sum(b * x)
+        w = csdl.sum(a * 2.0)  # uses the intermediate `a` outside x -> z
+        if do_compress:
+            op = compress(x, z, intermediates=intermediates)
+        g = csdl.derivative(z + w, x)
+        rec.stop()
+        results.append((w.value, g.value))
+    if intermediates == 'used':
+        assert op.extra_outputs == [a]
+    else:
+        assert set(op.extra_outputs) >= {a, b} and z not in op.extra_outputs
+    assert op.outputs[0] is z
+    np.testing.assert_allclose(results[1][0], results[0][0])
+    np.testing.assert_allclose(results[1][1], results[0][1], rtol=1e-12)
+
+
+def test_device_and_jit_kwargs_reach_jit_and_derivatives():
+    import jax
+    cpu = jax.devices('cpu')[0]
+    ref_rec, ref = build(False)
+    ref_g = csdl.derivative(ref['f'], ref['x'])
+    ref_rec.stop()
+
+    rec, v = build(True, device='cpu', jit_kwargs={'keep_unused': True})
+    g = csdl.derivative(v['f'], v['x'], loop=False)
+    rec.stop()
+    np.testing.assert_allclose(g.value, ref_g.value, rtol=1e-12)
+    from csdl_alpha.src.operations.operation_subclasses import SubgraphOperation as _S
+    vjp_ops = []
+    stack = [rec.active_graph]
+    while stack:
+        for n in stack.pop().node_table:
+            if isinstance(n, JaxCompressedOperation) and n.name.startswith('vjp'):
+                vjp_ops.append(n)
+            elif isinstance(n, _S):
+                stack.append(n.get_subgraph())
+    for op in [v['op']] + vjp_ops:
+        assert op.device == cpu and op.jit_kwargs == {'keep_unused': True}
+    assert vjp_ops
+
+
+def test_bad_options_raise_before_touching_graph():
+    rec = csdl.Recorder(inline=True)
+    rec.start()
+    x = csdl.Variable(value=np.ones(3))
+    z = csdl.sum(csdl.sin(x))
+    before = op_names(rec)
+    with pytest.raises(ValueError, match='No JAX device'):
+        compress(x, z, device='tpu')
+    with pytest.raises(ValueError, match='intermediates must be'):
+        compress(x, z, intermediates='some')
+    assert op_names(rec) == before
+    rec.stop()
+
+
+def test_output_independent_of_inputs_raises():
+    rec = csdl.Recorder(inline=True)
+    rec.start()
+    x = csdl.Variable(value=np.ones(3))
+    y = csdl.Variable(value=np.ones(3))
+    z = csdl.sum(y)
+    with pytest.raises(ValueError, match='does not affect'):  # CSDL's own check
+        compress(x, z)
+    rec.stop()
+
+
+def test_input_depending_on_input_raises():
+    rec = csdl.Recorder(inline=True)
+    rec.start()
+    x = csdl.Variable(value=np.ones(3))
+    y = csdl.sin(x)
+    z = csdl.sum(x * y)
+    with pytest.raises(ValueError, match='depends on another input'):
+        compress([x, y], z)
+    rec.stop()
+
+
+def test_derivative_is_one_compiled_node():
+    from csdl_alpha.src.graph.operation import Operation
+
+    def visible_ops(graph):
+        """Operations in `graph` and nested subgraphs, not looking inside JaxCompressedOperations."""
+        for node in graph.node_table:
+            if isinstance(node, Operation):
+                yield node
+                if isinstance(node, SubgraphOperation) and not isinstance(node, JaxCompressedOperation):
+                    yield from visible_ops(node.get_subgraph())
+
+    rec, v = build(True)
+    csdl.derivative(v['f'], v['x'])
+    rec.stop()
+    names = [op.name for op in visible_ops(rec.active_graph)]
+    assert names.count('vjp_jax_compressed') == 1
+    assert not {'sin', 'exp', 'matvec', 'norm'} & set(names)  # region and its VJP stay inside
+    assert isinstance(v['op'], SubgraphOperation)
+    assert 'sin' in {op.name for op in visible_ops(v['op'].get_subgraph())}
+
+
+def _implicit_model():
+    a = csdl.Variable(name='a', value=np.array([2.0, 3.0]))
+    x = csdl.ImplicitVariable(name='x', value=np.ones(2))
+    solver = csdl.nonlinear_solvers.Newton(print_status=False, tolerance=1e-12)
+    solver.add_state(x, x**3 + a * x - 10.0)
+    solver.run()
+    return a, csdl.sum(x**2 * a)
+
+
+def test_implicit_operation_inside_region():
+    # CSDL's implicit-function-theorem derivative is compiled into the VJP op.
+    results = []
+    for do_compress in [False, True]:
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        a, z = _implicit_model()
+        if do_compress:
+            compress(a, z)
+        g = csdl.derivative(z, a)
+        H = csdl.derivative(g, a)
+        rec.stop()
+        results.append((z.value, g.value, H.value))
+    for ref, got in zip(*results):
+        np.testing.assert_allclose(got, ref, rtol=1e-10, atol=1e-12)
+
+
+def test_implicit_operation_in_jax_simulator():
+    totals = []
+    for do_compress in [False, True]:
+        rec = csdl.Recorder(inline=False)
+        rec.start()
+        a, z = _implicit_model()
+        if do_compress:
+            compress(a, z)
+        rec.stop()
+        sim = csdl.experimental.JaxSimulator(rec, gpu=False, additional_inputs=[a], additional_outputs=[z])
+        sim[a] = np.array([1.5, 4.0])
+        sim.run()
+        totals.append((sim[z], sim.compute_totals()[z, a]))
+    np.testing.assert_allclose(totals[1][0], totals[0][0], rtol=1e-12)
+    np.testing.assert_allclose(totals[1][1], totals[0][1], rtol=1e-10)
+
+
+class CubePlusProduct(csdl.CustomExplicitOperation):
+    """f = x**3 + x*y, elementwise, with hand-written derivatives."""
+
+    def evaluate(self, x, y):
+        self.declare_input('x', x)
+        self.declare_input('y', y)
+        f = self.create_output('f', x.shape)
+        self.declare_derivative_parameters('f', 'x')
+        self.declare_derivative_parameters('f', 'y')
+        return f
+
+    def compute(self, input_vals, output_vals):
+        output_vals['f'] = input_vals['x']**3 + input_vals['x'] * input_vals['y']
+
+    def compute_derivatives(self, input_vals, output_vals, derivatives):
+        derivatives['f', 'x'] = np.diag(3 * input_vals['x']**2 + input_vals['y'])
+        derivatives['f', 'y'] = np.diag(input_vals['x'])
+
+
+def _custom_model(x, y):
+    u = csdl.sin(x) * 2.0
+    f = CubePlusProduct().evaluate(u, y)
+    return csdl.sum(csdl.exp(0.1 * f) * u)
+
+
+def _build_custom(do_compress, inline, **compress_kwargs):
+    rec = csdl.Recorder(inline=inline)
+    rec.start()
+    x = csdl.Variable(name='x', value=np.array([0.3, -0.4, 0.9]))
+    y = csdl.Variable(name='y', value=np.array([1.0, 0.5, -2.0]))
+    z = _custom_model(x, y)
+    op = compress([x, y], z, **compress_kwargs) if do_compress else None
+    f = 3.0 * z
+    return rec, x, y, f, op
+
+
+@pytest.mark.parametrize('loop', [True, False])
+def test_custom_operation_uses_csdl_derivatives(loop):
+    ref_rec, rx_, ry_, rf, _ = _build_custom(False, True)
+    ref_d = csdl.derivative(rf, [rx_, ry_], loop=loop)
+    ref_rec.stop()
+
+    rec, x, y, f, op = _build_custom(True, True)
+    d = csdl.derivative(f, [x, y], loop=loop)
+    rec.stop()
+
+    np.testing.assert_allclose(f.value, rf.value)
+    np.testing.assert_allclose(d[x].value, ref_d[rx_].value, rtol=1e-12)
+    np.testing.assert_allclose(d[y].value, ref_d[ry_].value, rtol=1e-12)
+
+
+def test_custom_operation_in_simulators():
+    ref_rec, rx_, ry_, rf, _ = _build_custom(False, False)
+    ref_rec.stop()
+    ref_sim = csdl.experimental.PySimulator(ref_rec)
+    ref_sim.run()
+    ref_d = ref_sim.compute_totals(rf, [rx_, ry_])
+
+    rec, x, y, f, _ = _build_custom(True, False)
+    rec.stop()
+    for sim in [csdl.experimental.PySimulator(rec),
+                csdl.experimental.JaxSimulator(rec, gpu=False, additional_inputs=[x, y], additional_outputs=[f])]:
+        sim.run()
+        np.testing.assert_allclose(sim[f], ref_sim[rf])
+        d = sim.compute_totals(f, [x, y]) if isinstance(sim, csdl.experimental.PySimulator) else sim.compute_totals()
+        np.testing.assert_allclose(d[f, x], ref_d[rf, rx_], rtol=1e-12)
+        np.testing.assert_allclose(d[f, y], ref_d[rf, ry_], rtol=1e-12)
+
+
+def _jax_sim_results(rec, x, c, f, loop):
+    sim = csdl.experimental.JaxSimulator(
+        rec, gpu=False, additional_inputs=[x, c], additional_outputs=[f],
+        derivatives_kwargs={'loop': loop})
+    sim[x] = np.array([0.3, -0.1, 0.9, 0.2])
+    sim.run()
+    d = sim.compute_totals()
+    return sim[f], d[f, x], d[f, c]
+
+
+@pytest.mark.parametrize('compile_separately', [True, False])
+@pytest.mark.parametrize('loop', [True, False])  # loop=False vmaps the VJP op
+def test_both_compile_modes_under_jax_simulator(loop, compile_separately):
+    ref_rec, ref = build(False, inline=False)
+    ref_rec.stop()
+    expected = _jax_sim_results(ref_rec, ref['x'], ref['c'], ref['f'], loop)
+
+    rec, v = build(True, inline=False, compile_separately=compile_separately)
+    rec.stop()
+    got = _jax_sim_results(rec, v['x'], v['c'], v['f'], loop)
+    for e, g in zip(expected, got):
+        np.testing.assert_allclose(g, e, rtol=1e-12)
+
+
+def test_compile_separately_keeps_region_out_of_enclosing_program():
+    import jax
+    from csdl_alpha.backends.jax.graph_to_jax import create_jax_function
+
+    def primitives(compile_separately):
+        rec, v = build(True, inline=False, compile_separately=compile_separately)
+        rec.stop()
+        fn = create_jax_function(rec.active_graph, [v['f']], [v['x'], v['y'], v['c']])
+        jaxpr = jax.make_jaxpr(fn)(*(jax.numpy.asarray(v[k].value) for k in 'xyc'))
+        return [eqn.primitive.name for eqn in jaxpr.eqns]
+
+    assert 'pjit' in primitives(False) or 'jit' in primitives(False)
+    separate = primitives(True)
+    assert 'pure_callback' in separate and not {'pjit', 'jit'} & set(separate)
+
+
+def test_compile_separately_nested_and_with_custom_and_implicit_ops():
+    results = []
+    for do_compress in [False, True]:
+        rec = csdl.Recorder(inline=False)
+        rec.start()
+        a, z_implicit = _implicit_model()
+        x = csdl.Variable(name='x', value=np.array([0.3, -0.4]))
+        z_custom = _custom_model(x, a)
+        if do_compress:
+            compress([a, x], z_custom)                                   # inner, separate (default)
+        f = z_implicit * z_custom
+        if do_compress:
+            compress([a, x], f, name='outer', compile_separately=False)  # outer, fused, contains the inner
+        rec.stop()
+        sim = csdl.experimental.JaxSimulator(rec, gpu=False, additional_inputs=[a, x], additional_outputs=[f])
+        sim[a] = np.array([1.5, 4.0])
+        sim.run()
+        d = sim.compute_totals()
+        results.append((sim[f], d[f, a], d[f, x]))
+    for ref, got in zip(*results):
+        np.testing.assert_allclose(got, ref, rtol=1e-10)

From 8dc294123e7127f961a891c98ccb67a50e6d545b Mon Sep 17 00:00:00 2001
From: Michael Warner <hgxe14@gmail.com>
Date: Thu, 8 Oct 2026 16:24:14 -0700
Subject: [PATCH 2/4] Test compress inside frange and enter_loop bodies

Compressing a region inside a loop body (frange or enter_loop, with and
without the loop index, fused and separate) matches the uncompressed loop
inline and under JaxSimulator, and the body holds one compressed operation.
TestJaxCompress gains the same case for the backend matrix.
---
 .../src/operations/compress_operations.py     | 37 ++++++++++
 .../src/operations/tests/test_compress.py     | 68 +++++++++++++++++++
 2 files changed, 105 insertions(+)

diff --git a/csdl_alpha/src/operations/compress_operations.py b/csdl_alpha/src/operations/compress_operations.py
index 8884c60..7676aee 100644
--- a/csdl_alpha/src/operations/compress_operations.py
+++ b/csdl_alpha/src/operations/compress_operations.py
@@ -635,6 +635,43 @@ def test_loop_and_nested(self):
         z_np = np.sum(sum(np.sin(x_val * (i + 1.0)) for i in range(3)) * x_val)
         self.run_tests([csdl_tests.TestingPair(f, np.array([np.exp(0.1 * z_np)]))], verify_derivatives=True)
 
+    def test_compress_inside_loop_body(self):
+        import csdl_alpha as csdl
+        x_val, p_val = np.array([0.2, 0.5, 0.9]), np.array([1.3])
+
+        def step(h, i, p, lib):
+            return lib.tanh(lib.sin(h * p) + (i + 1.0) * 0.3 * h) * 0.9 + h * 0.1
+
+        h_np = x_val
+        for i in range(4):
+            h_np = step(h_np, i, p_val, np)
+        f_np = np.array([np.sum(h_np**2)])
+
+        for loop_kind in ['frange', 'enter_loop']:
+            for compile_separately in [True, False]:
+                self.prep()
+                x = csdl.Variable(name='x', value=x_val)
+                p = csdl.Variable(name='p', value=p_val)
+                if loop_kind == 'frange':
+                    h = x
+                    for i in csdl.frange(4):
+                        h_new = step(h, i, p, csdl)
+                        compress(h, h_new, compile_separately=compile_separately)
+                        h = h_new
+                else:
+                    with csdl.experimental.enter_loop(vals=[list(range(4))]) as loop_builder:
+                        i = loop_builder.get_loop_indices()
+                        h0 = loop_builder.initialize_feedback(x)
+                        h1 = step(h0, i, p, csdl)
+                        compress(h0, h1, compile_separately=compile_separately)
+                        loop_builder.finalize_feedback(h0, h1)
+                    h = loop_builder.add_output(h1)
+                    loop_builder.finalize()
+                f = csdl.sum(h**2)
+                self.run_tests(
+                    [csdl_tests.TestingPair(f, f_np, tag=f'{loop_kind} separate={compile_separately}')],
+                    verify_derivatives=True)
+
 
 if __name__ == '__main__':
     test = TestCompressOp()
diff --git a/csdl_alpha/src/operations/tests/test_compress.py b/csdl_alpha/src/operations/tests/test_compress.py
index a118c49..44611ec 100644
--- a/csdl_alpha/src/operations/tests/test_compress.py
+++ b/csdl_alpha/src/operations/tests/test_compress.py
@@ -500,3 +500,71 @@ def test_compile_separately_nested_and_with_custom_and_implicit_ops():
         results.append((sim[f], d[f, a], d[f, x]))
     for ref, got in zip(*results):
         np.testing.assert_allclose(got, ref, rtol=1e-10)
+
+
+def _loop_model(compress_mode, loop_kind, use_index, inline):
+    """Four iterations of h <- 0.9 tanh(sin(p h) + s h) + 0.1 h, optionally compressing the body's region."""
+    rec = csdl.Recorder(inline=inline)
+    rec.start()
+    x = csdl.Variable(name='x', value=np.array([0.2, 0.5, 0.9]))
+    p = csdl.Variable(name='p', value=np.array([1.3]))  # read from outside the loop
+
+    def body(h, i):
+        scale = (i + 1.0) * 0.3 if use_index else 0.3
+        h_new = csdl.tanh(csdl.sin(h * p) + scale * h) * 0.9 + h * 0.1
+        if compress_mode is not None:
+            compress(h, h_new, compile_separately=(compress_mode == 'separate'))
+        return h_new
+
+    if loop_kind == 'frange':
+        h = x
+        for i in csdl.frange(4):
+            h = body(h, i)
+    else:
+        with csdl.experimental.enter_loop(vals=[list(range(4))]) as loop_builder:
+            i = loop_builder.get_loop_indices()
+            h0 = loop_builder.initialize_feedback(x)
+            h1 = body(h0, i)
+            loop_builder.finalize_feedback(h0, h1)
+        h = loop_builder.add_output(h1)
+        loop_builder.finalize()
+    return rec, x, p, csdl.sum(h**2)
+
+
+@pytest.mark.parametrize('compress_mode', ['fused', 'separate'])
+@pytest.mark.parametrize('use_index', [False, True])
+@pytest.mark.parametrize('loop_kind', ['frange', 'enter_loop'])
+def test_compress_inside_loop_body(loop_kind, use_index, compress_mode):
+    from csdl_alpha.src.graph.operation import Operation
+
+    def inline_results(mode):
+        rec, x, p, f = _loop_model(mode, loop_kind, use_index, inline=True)
+        results = [f.value]
+        for loop in [True, False]:
+            d = csdl.derivative(f, [x, p], loop=loop)
+            results += [d[x].value, d[p].value]
+        rec.stop()
+        return rec, results
+
+    def jax_sim_results(mode):
+        rec, x, p, f = _loop_model(mode, loop_kind, use_index, inline=False)
+        rec.stop()
+        sim = csdl.experimental.JaxSimulator(rec, gpu=False, additional_inputs=[x, p], additional_outputs=[f])
+        sim[x] = np.array([0.4, -0.3, 1.1])
+        sim.run()
+        d = sim.compute_totals()
+        return [sim[f], d[f, x], d[f, p]]
+
+    _, expected = inline_results(None)
+    rec, got = inline_results(compress_mode)
+    for g, e in zip(got, expected):
+        np.testing.assert_allclose(g, e, rtol=1e-10, atol=1e-12)
+    for g, e in zip(jax_sim_results(compress_mode), jax_sim_results(None)):
+        np.testing.assert_allclose(g, e, rtol=1e-10, atol=1e-12)
+
+    # The loop body holds one compressed operation in place of the sin/tanh chain.
+    loops = [n for n in rec.active_graph.node_table
+             if isinstance(n, SubgraphOperation) and not isinstance(n, JaxCompressedOperation)]
+    body_names = [n.name for n in loops[0].get_subgraph().node_table if isinstance(n, Operation)]
+    assert body_names.count('jax_compressed') == 1
+    assert not {'sin', 'tanh'} & set(body_names)

From e0edc547ef0dcd3a8cb127b616d841aae79fc93d Mon Sep 17 00:00:00 2001
From: Michael Warner <hgxe14@gmail.com>
Date: Thu, 8 Oct 2026 18:45:24 -0700
Subject: [PATCH 3/4] Share compiled functions between identical compressed
 regions

- share_compiled (default True): a new JaxCompressedOperation reuses the
  compiled function of an earlier identical one in the recorder
  (Recorder.compressed_operations). Regions are identical when a structural
  walk pairs their operations port by port (class, shapes, wiring, constant
  values) and one JAX trace of each region, in paired order, gives the same
  jaxpr. Inputs may be in a different order.
- share_with=op: reuse op's compiled function explicitly; the region is
  checked and a mismatch raises before the graph changes.
- find_repeats: also compress every other copy of the region in the active
  graph with the same function (greedy in topological order; op.repeats).
- assume_custom_ops_match: JAX cannot see what a callback computes, so
  regions containing custom operations only match when asserted (share_with
  asserts it for that pair).
- Operations sharing a function form a group; derivative operations are
  shared within the group by (seeded outputs, differentiated inputs),
  without matching, since they differentiate the same function.
---
 .../src/operations/compress_operations.py     | 530 +++++++++++++++++-
 .../src/operations/tests/test_compress.py     | 363 ++++++++++++
 csdl_alpha/src/recorder.py                    |   5 +
 3 files changed, 882 insertions(+), 16 deletions(-)

diff --git a/csdl_alpha/src/operations/compress_operations.py b/csdl_alpha/src/operations/compress_operations.py
index 7676aee..37c3700 100644
--- a/csdl_alpha/src/operations/compress_operations.py
+++ b/csdl_alpha/src/operations/compress_operations.py
@@ -166,8 +166,23 @@ class JaxCompressedOperation(CompressedOperation):
         boundary; under ``vmap`` the callback runs once per batch element. With
         False, the operation is traced into the enclosing program instead.
 
-    Derivative operations inherit ``device``, ``jit_kwargs``, and
-    ``compile_separately``.
+    share_compiled : bool
+        Reuse the compiled function of an earlier ``JaxCompressedOperation`` in
+        this recorder whose subgraph is identical (same operations, parameters,
+        shapes, wiring, and compiled-in constants), by default True. The
+        earlier operation is ``shared_from``.
+    assume_custom_ops_match : bool
+        Treat operations that JAX runs through a Python callback (custom
+        operations, operations without a JAX implementation) as identical when
+        their class, shapes, and JAX trace match, by default False. JAX cannot
+        see the Python state such an operation computes with (for example a
+        parameter stored on the instance), so by default subgraphs containing
+        one are never shared. Set True only when you know those operations
+        compute the same function.
+
+    Derivative operations inherit ``device``, ``jit_kwargs``,
+    ``compile_separately``, ``share_compiled``, and
+    ``assume_custom_ops_match``.
     """
 
     def __init__(
@@ -179,6 +194,8 @@ def __init__(
             device = None,
             jit_kwargs = None,
             compile_separately = True,
+            share_compiled = True,
+            assume_custom_ops_match = False,
         ):
         import jax
         jax.config.update("jax_enable_x64", True)
@@ -186,15 +203,32 @@ def __init__(
         self.device = _resolve_device(device)
         self.jit_kwargs = dict(jit_kwargs or {})
         self.compile_separately = compile_separately
+        self.share_compiled = share_compiled
+        self.assume_custom_ops_match = assume_custom_ops_match
+        # The operation whose compiled function this one reuses, and, for each
+        # of that operation's inputs, the position of the matching input here.
+        self.shared_from = None
+        self._input_permutation = None
+        self._share_group = _ShareGroup(self)
+        if share_compiled:
+            _share_with_identical(self)
 
     @property
     def jit_fn(self):
         """The subgraph as a ``jax.jit`` function, built on first use."""
         if self.jax_function is None:
             import jax
-            from csdl_alpha.backends.jax.graph_to_jax import create_jax_function
-            fn = create_jax_function(self.get_subgraph(), self.outputs, self.inputs)
-            self.jax_function = jax.jit(fn, **self.jit_kwargs)
+            if self.shared_from is not None:
+                shared_fn = self.shared_from.jit_fn
+                permutation = self._input_permutation
+                if permutation == list(range(len(permutation))):
+                    self.jax_function = shared_fn
+                else:
+                    self.jax_function = lambda *args: shared_fn(*(args[j] for j in permutation))
+            else:
+                from csdl_alpha.backends.jax.graph_to_jax import create_jax_function
+                fn = create_jax_function(self.get_subgraph(), self.outputs, self.inputs)
+                self.jax_function = jax.jit(fn, **self.jit_kwargs)
         return self.jax_function
 
     def compute_jax(self, *args):
@@ -225,11 +259,27 @@ def evaluate_vjp(self, cotangents, *inputs_outputs):
 
         inputs = inputs_outputs[:self.num_inputs]
         outputs = inputs_outputs[self.num_inputs:]
+        if self._input_permutation is not None:
+            # List inputs in the order of the operation that owns the compiled
+            # function, so derivative operations of operations that share it
+            # come out in one order and can share too.
+            inputs = [inputs[j] for j in self._input_permutation]
         # Outputs with a None cotangent contribute nothing; leave them out.
-        seeds = [(y, cotangents[y]) for y in outputs if cotangents.check(y) and cotangents[y] is not None]
-        wrts = [x for x in inputs if cotangents.check(x)]
-        if not seeds or not wrts:
+        seeded = [j for j, y in enumerate(outputs) if cotangents.check(y) and cotangents[y] is not None]
+        differentiated = [i for i, x in enumerate(inputs) if cotangents.check(x)]
+        if not seeded or not differentiated:
             return
+        seeds = [(outputs[j], cotangents[outputs[j]]) for j in seeded]
+        wrts = [inputs[i] for i in differentiated]
+
+        # Every operation in this group computes the same function, so a
+        # derivative operation another member made for the same seeded outputs
+        # and differentiated inputs computes the same function too: share its
+        # compiled function without checking. (Repeated variables among the
+        # inputs and seeds would change the derivative operation's inputs.)
+        vjp_inputs = list(inputs) + [seed for _, seed in seeds]
+        pattern = (tuple(seeded), tuple(differentiated)) if len(set(vjp_inputs)) == len(vjp_inputs) else None
+        shared_vjp = self._share_group.derivatives.get(pattern) if pattern is not None else None
 
         # Record CSDL's reverse mode into a copy of the subgraph, so the subgraph
         # itself stays a clean forward graph for later derivative calls. The
@@ -258,8 +308,14 @@ def evaluate_vjp(self, cotangents, *inputs_outputs):
             recorder._add_node(grad)
         # The copy also holds forward operations the VJP does not need; jax.jit drops them.
         vjp_op = JaxCompressedOperation(
-            work, list(inputs) + [seed for _, seed in seeds], grads, name=f'vjp_{self.name}',
-            device=self.device, jit_kwargs=self.jit_kwargs, compile_separately=self.compile_separately)
+            work, vjp_inputs, grads, name=f'vjp_{self.name}',
+            device=self.device, jit_kwargs=self.jit_kwargs, compile_separately=self.compile_separately,
+            share_compiled=False, assume_custom_ops_match=self.assume_custom_ops_match)
+        vjp_op.share_compiled = self.share_compiled
+        if shared_vjp is not None:
+            _share_explicitly(vjp_op, shared_vjp, list(range(len(vjp_inputs))))
+        elif pattern is not None:
+            self._share_group.derivatives[pattern] = vjp_op
         vjp_op.finalize_and_return_outputs()
         for x, grad in zip(wrts, grads):
             cotangents.accumulate(x, grad)
@@ -274,6 +330,10 @@ def compress(
         device = None,
         jit_kwargs = None,
         compile_separately = True,
+        share_compiled = True,
+        find_repeats = False,
+        assume_custom_ops_match = False,
+        share_with = None,
     ):
     """Replace the operations between ``inputs`` and ``outputs`` with one JAX-compiled operation.
 
@@ -321,14 +381,51 @@ def compress(
         program, by default True. Set False to let XLA compile and optimize
         the region together with the rest of the program. See
         :class:`JaxCompressedOperation`.
+    share_compiled : bool
+        Reuse the compiled function of an earlier compressed operation in this
+        recorder whose region is identical, by default True. Its derivative
+        operations share in the same way.
+    find_repeats : bool
+        Also find every other copy of this region in the active graph (same
+        operations, parameters, shapes, wiring, and constants, reading any
+        inputs) and compress each one with the same compiled function, by
+        default False. Copies that have an intermediate used outside them or
+        would create a cycle are skipped. The copies are ``op.repeats``.
+
+        Copies can overlap (in a chain ``x1 = sin(x0)``, ``x2 = sin(x1)``, ...,
+        a two-``sin`` region matches at every offset). The region given here
+        is always compressed as given, and copies are then taken greedily in
+        topological order of their outputs, skipping any that would reuse an
+        operation already compressed. This gives the most copies when
+        overlapping copies line up like a chain, but not necessarily in
+        branching graphs. To choose the copies exactly, compress each one and
+        pass ``share_with``.
+    assume_custom_ops_match : bool
+        Let regions containing custom operations (or other operations JAX runs
+        through a Python callback) count as identical when the operations'
+        class, shapes, and JAX trace match, by default False. JAX cannot see
+        the Python state a custom operation computes with, so only set this
+        when you know the custom operations in matching regions compute the
+        same function.
+    share_with : JaxCompressedOperation, optional
+        An earlier compressed operation whose compiled function this region
+        should reuse. The region is checked against it (same operations,
+        parameters, shapes, wiring, and constants; ``outputs`` in the same
+        order; inputs in any order) and a mismatch raises before the graph is
+        changed. Custom operations of the same class count as matching, since
+        naming ``share_with`` asserts that the regions compute the same
+        function. Derivative operations share as well. ``jit_kwargs`` are
+        taken from ``share_with``.
 
     Returns
     -------
     JaxCompressedOperation
         ``op.inputs`` is ``inputs`` followed by ``op.extra_inputs``,
         ``op.outputs`` is ``outputs`` followed by ``op.extra_outputs``
-        (intermediates added by ``intermediates``), and ``op.get_subgraph()``
-        is the region.
+        (intermediates added by ``intermediates``), ``op.get_subgraph()``
+        is the region, ``op.shared_from`` is the operation whose compiled
+        function it reuses (or None), and ``op.repeats`` lists the copies
+        compressed by ``find_repeats``.
 
     Examples
     --------
@@ -342,11 +439,16 @@ def compress(
     array([[0.39866767, 0.78936529]])
     """
     import csdl_alpha as csdl
-    from csdl_alpha.src.graph.variable import Constant
     from csdl_alpha.utils.inputs import listify_variables
 
     if intermediates not in ('none', 'used', 'all'):
         raise ValueError(f"intermediates must be 'none', 'used', or 'all', not {intermediates!r}.")
+    if share_with is not None:
+        if not isinstance(share_with, JaxCompressedOperation):
+            raise TypeError(f"share_with must be a JaxCompressedOperation, not {type(share_with).__name__}.")
+        if jit_kwargs is not None and dict(jit_kwargs) != share_with.jit_kwargs:
+            raise ValueError("jit_kwargs differ from share_with's; the shared compiled function uses share_with's.")
+        jit_kwargs = share_with.jit_kwargs
     device = _resolve_device(device)
     inputs = listify_variables(inputs)
     outputs = listify_variables(outputs)
@@ -377,10 +479,16 @@ def compress(
         for var in op.inputs:
             if var not in internal and var not in input_set and var not in seen:
                 seen.add(var)
-                is_literal = type(var) is Constant and var.value is not None
-                (constants if is_literal else extra_inputs).append(var)
+                (constants if _is_literal(var) else extra_inputs).append(var)
     op_inputs = inputs + extra_inputs
 
+    if share_with is not None:
+        share_permutation = _match(share_with, graph, op_inputs, outputs, set(ops), assume_custom_ops_match=True)
+        if share_permutation is None:
+            raise ValueError(
+                f"The region does not match the region of share_with ({share_with.info()}): they differ in "
+                "operations, parameters, shapes, wiring, constants, or the order of `outputs`.")
+
     # Move the region's operations into their own graph (this deletes them from
     # `graph` and keeps the variables there), then drop the intermediates and
     # any constants nothing else uses. The outputs stay; the new operation
@@ -395,11 +503,16 @@ def compress(
 
     compressed = JaxCompressedOperation(
         region, op_inputs, outputs, name=name, device=device,
-        jit_kwargs=jit_kwargs, compile_separately=compile_separately)
+        jit_kwargs=jit_kwargs, compile_separately=compile_separately,
+        share_compiled=share_compiled and share_with is None, assume_custom_ops_match=assume_custom_ops_match)
+    if share_with is not None:
+        compressed.share_compiled = share_compiled  # for its derivative operations
+        _share_explicitly(compressed, share_with, share_permutation)
     compressed.extra_inputs = extra_inputs
     compressed.extra_outputs = extra_outputs
     # Skip inline evaluation: the outputs already hold the values the region computed.
     compressed.finalize_and_return_outputs(skip_inline=True)
+    compressed.repeats = _compress_repeats(compressed, graph, len(inputs)) if find_repeats else []
     return compressed
 
 
@@ -482,6 +595,357 @@ def producers_of_inputs(op_index):
             stack.extend(producers_of_inputs(candidate))
 
 
+# ---- Sharing compiled functions between identical regions ----------------------
+#
+# Two regions are identical when a one-to-one map pairs their operations and
+# variables such that paired operations are of the same class with the same
+# shapes, every operation input is wired from the paired producer output at the
+# same port, and compiled-in constants have equal values (the structural walk),
+# and the two regions, traced with JAX operation by operation in paired order,
+# give the same jaxpr (which checks parameters the walk cannot see, such as
+# reshape shapes or exponents). Template inputs may map to any variables.
+# Because operation inputs are ordered and each variable has one producer, the
+# walk goes back from the outputs port by port, with no search except where an
+# operation is reachable only forward from a variable (see _complete).
+#
+# Operations that share a compiled function form a _ShareGroup. Their
+# derivative operations are shared within the group without matching, since
+# they differentiate the same function.
+
+
+class _ShareGroup:
+    """Operations that compute one function, and the derivative operations made
+    for them, keyed by (seeded outputs, differentiated inputs) in the order of
+    the operation that owns the compiled function."""
+
+    def __init__(self, owner):
+        self.members = [owner]
+        self.derivatives = {}
+
+
+def _is_literal(var):
+    from csdl_alpha.src.graph.variable import Constant
+    return type(var) is Constant and var.value is not None
+
+
+def _registry_key(op):
+    num_ops = sum(isinstance(node, Operation) for node in op.get_subgraph().node_table)
+    return (
+        tuple(sorted(var.shape for var in op.inputs)),
+        tuple(var.shape for var in op.outputs),
+        num_ops,
+        repr(sorted(op.jit_kwargs.items())),
+    )
+
+
+def _root(op):
+    return op.shared_from if op.shared_from is not None else op
+
+
+def _share_explicitly(op, template, permutation):
+    """Make ``op`` use ``template``'s compiled function; ``permutation[i]`` is the
+    position in ``op.inputs`` of ``template``'s input ``i``."""
+    if template.shared_from is not None:
+        # Compose with the template's own mapping onto the function's owner.
+        permutation = [permutation[j] for j in template._input_permutation]
+    op.shared_from = _root(template)
+    op._input_permutation = permutation
+    template._share_group.members.append(op)
+    op._share_group = template._share_group
+
+
+def _share_with_identical(op):
+    """Point ``op`` at the compiled function of an identical earlier operation, or register it."""
+    entries = op.recorder.compressed_operations.setdefault(_registry_key(op), [])
+    for template in entries:
+        permutation = _match(template, op.get_subgraph(), op.inputs, op.outputs, None, op.assume_custom_ops_match)
+        if permutation is not None:
+            _share_explicitly(op, template, permutation)
+            return
+    entries.append(op)
+
+
+def _same_structure(t_op, c_op, ctx):
+    """Whether two operations can pair in the walk (class and shapes); their
+    parameters are compared later through the region trace."""
+    if type(t_op) is not type(c_op):
+        return False
+    if len(t_op.inputs) != len(c_op.inputs) or len(t_op.outputs) != len(c_op.outputs):
+        return False
+    if any(a.shape != b.shape for a, b in zip(t_op.inputs + t_op.outputs, c_op.inputs + c_op.outputs)):
+        return False
+    if isinstance(t_op, JaxCompressedOperation):
+        # Nested compressed operations sharing one compiled function, with the
+        # same input order, are identical: trace them as a placeholder.
+        if _root(t_op) is _root(c_op) and t_op._input_permutation == c_op._input_permutation:
+            ctx.placeholders.add(t_op)
+    return True
+
+
+class _MatchState:
+    def __init__(self):
+        self.var_map, self.op_map, self.used_vars, self.used_ops = {}, {}, set(), set()
+
+    def copy(self):
+        new = _MatchState()
+        new.var_map, new.op_map = dict(self.var_map), dict(self.op_map)
+        new.used_vars, new.used_ops = set(self.used_vars), set(self.used_ops)
+        return new
+
+
+class _MatchContext:
+    def __init__(self, template, c_graph, input_ok, leaf_ok, op_ok, assume_custom_ops_match):
+        self.template = template
+        self.t_graph = template.get_subgraph()
+        self.t_inputs = set(template.inputs)
+        self.t_ops = _live_ops(template)
+        self.c_graph, self.input_ok, self.leaf_ok, self.op_ok = c_graph, input_ok, leaf_ok, op_ok
+        self.assume = assume_custom_ops_match
+        self.placeholders = set()  # template operations traced as placeholders
+
+
+def _live_ops(op):
+    """Operations in ``op``'s subgraph that its outputs depend on, in topological order."""
+    if '_compress_live_ops' not in op.__dict__:
+        import rustworkx as rx
+        graph = op.get_subgraph()
+        indices = set()
+        for var in op.outputs:
+            indices |= rx.ancestors(graph.rxgraph, graph.node_table[var])
+        op._compress_live_ops = [graph.rxgraph[i] for i in rx.topological_sort(graph.rxgraph)
+                                 if i in indices and isinstance(graph.rxgraph[i], Operation)]
+    return op._compress_live_ops
+
+
+def _walk(state, pairs, ctx):
+    """Extend ``state`` by pairing template and candidate variables back to their producers; None on mismatch."""
+    stack = list(pairs)
+    while stack:
+        t, c = stack.pop()
+        if t in state.var_map:
+            if state.var_map[t] is not c:
+                return None
+            continue
+        if c in state.used_vars or t.shape != c.shape:
+            return None
+        state.var_map[t] = c
+        state.used_vars.add(c)
+        if t in ctx.t_inputs:
+            if not ctx.input_ok(c):
+                return None
+            continue
+        t_producers = ctx.t_graph.predecessors(t)
+        if not t_producers:  # a value compiled into the template's function
+            if not ctx.leaf_ok(c) or not np.array_equal(t.value, c.value):
+                return None
+            continue
+        t_op = t_producers[0]
+        c_producers = ctx.c_graph.predecessors(c) if c in ctx.c_graph.node_table else []
+        if not c_producers:
+            return None
+        c_op = c_producers[0]
+        if t_op in state.op_map:
+            if state.op_map[t_op] is not c_op:
+                return None
+            continue
+        if c_op in state.used_ops or not ctx.op_ok(c_op) or not _same_structure(t_op, c_op, ctx):
+            return None
+        state.op_map[t_op] = c_op
+        state.used_ops.add(c_op)
+        stack.extend(zip(t_op.outputs, c_op.outputs))
+        stack.extend(zip(t_op.inputs, c_op.inputs))
+    return state
+
+
+def _complete(state, ctx):
+    """Map template operations the backward walk did not reach, by trying the
+    consumers of an already-mapped variable at the same input port."""
+    missing = [op for op in ctx.t_ops if op not in state.op_map]
+    if not missing:
+        return state
+    for t_op in missing:
+        bound = [(port, var) for port, var in enumerate(t_op.inputs) if var in state.var_map]
+        if bound:
+            break
+    else:
+        return None
+    port, t_var = bound[0]
+    c_var = state.var_map[t_var]
+    graph = ctx.c_graph
+    for c_op in graph.rxgraph.successors(graph.node_table[c_var]):
+        if len(c_op.inputs) > port and c_op.inputs[port] is c_var and type(c_op) is type(t_op):
+            trial = _walk(state.copy(), list(zip(t_op.outputs, c_op.outputs)), ctx)
+            if trial is not None and trial.op_map.get(t_op) is c_op:
+                result = _complete(trial, ctx)
+                if result is not None:
+                    return result
+    return None
+
+
+def _jax_outputs(op, args):
+    """``op``'s outputs as JAX arrays, as ``create_jax_function`` evaluates it."""
+    from csdl_alpha.src.operations.loops.new_loop.new_loop import NewLoop
+    if isinstance(op, JaxCompressedOperation):
+        return list(op.jit_fn(*args))  # its own trace, not the callback of compile_separately
+    if isinstance(op, NewLoop):
+        outputs = {var: None for var in op.outputs}
+        op.evaluate_jax(dict(zip(op.inputs, args)), outputs=outputs)
+        return [outputs[var] for var in op.outputs]
+    outs = op.compute_jax(*args)
+    return list(outs) if isinstance(outs, (tuple, list)) else [outs]
+
+
+def _trace_region(ops, inputs, outputs, placeholders):
+    """``(jaxpr text, constants)`` of running ``ops`` in the given order, or None if it cannot be traced."""
+    import jax
+    import jax.numpy as jnp
+
+    def region(*args):
+        env = dict(zip(inputs, args))
+        for op in ops:
+            for var in op.inputs:
+                if var not in env:  # a value compiled in
+                    env[var] = jnp.asarray(var.value)
+            if op in placeholders:
+                outs = [jnp.zeros(var.shape) for var in op.outputs]
+            else:
+                outs = _jax_outputs(op, [env[var] for var in op.inputs])
+            for var, out in zip(op.outputs, outs):
+                env[var] = jnp.reshape(out, var.shape)
+        return [env[var] for var in outputs]
+
+    try:
+        closed = jax.make_jaxpr(region)(*(jax.ShapeDtypeStruct(var.shape, np.float64) for var in inputs))
+    except Exception:
+        return None
+    return str(closed.jaxpr), tuple((np.shape(c), np.asarray(c).tobytes()) for c in closed.consts)
+
+
+def _traces_match(state, ctx, c_inputs):
+    """Whether the candidate paired in ``state`` traces like the template;
+    ``c_inputs`` are the candidate inputs in the template's input order."""
+    template = ctx.template
+    cache = template.__dict__.setdefault('_compress_traces', {})
+    key = frozenset(id(op) for op in ctx.placeholders)
+    if key not in cache:
+        cache[key] = _trace_region(ctx.t_ops, template.inputs, template.outputs, ctx.placeholders)
+    t_trace = cache[key]
+    c_ops = [state.op_map[op] for op in ctx.t_ops]
+    c_placeholders = {state.op_map[op] for op in ctx.placeholders}
+    c_outputs = [state.var_map[var] for var in template.outputs]
+    c_trace = _trace_region(c_ops, c_inputs, c_outputs, c_placeholders)
+    if t_trace is None or c_trace is None or t_trace != c_trace:
+        return False
+    # JAX cannot see what a callback computes, so identical traces do not
+    # imply identical regions unless the user says so.
+    return ctx.assume or 'callback' not in t_trace[0]
+
+
+def _match(template, c_graph, c_inputs, c_outputs, c_ops, assume_custom_ops_match):
+    """How a candidate region computes ``template``'s function: for each template
+    input, the position of the matching candidate input; None if they differ.
+
+    The candidate region is ``c_ops`` (all operations of ``c_graph`` if None)
+    with inputs ``c_inputs`` and outputs ``c_outputs``, which match the
+    template's outputs in order.
+    """
+    if len(c_inputs) != len(template.inputs) or len(c_outputs) != len(template.outputs):
+        return None
+    c_input_set = set(c_inputs)
+    ctx = _MatchContext(
+        template, c_graph,
+        input_ok=lambda var: var in c_input_set,
+        leaf_ok=lambda var: var not in c_input_set and c_graph.in_degree(var) == 0,
+        op_ok=(lambda op: True) if c_ops is None else (lambda op: op in c_ops),
+        assume_custom_ops_match=assume_custom_ops_match,
+    )
+    state = _walk(_MatchState(), list(zip(template.outputs, c_outputs)), ctx)
+    state = _complete(state, ctx) if state is not None else None
+    if state is None:
+        return None
+    # Inputs the function reads are mapped exactly; inputs it ignores pair up in order.
+    position = {var: j for j, var in enumerate(c_inputs)}
+    unused = [j for j, var in enumerate(c_inputs) if var not in state.used_vars]
+    permutation = []
+    for t_var in template.inputs:
+        if t_var in state.var_map:
+            permutation.append(position[state.var_map[t_var]])
+        elif unused and c_inputs[unused[0]].shape == t_var.shape:
+            permutation.append(unused.pop(0))
+        else:
+            return None
+    if unused or not _traces_match(state, ctx, [c_inputs[j] for j in permutation]):
+        return None
+    return permutation
+
+
+def _compress_repeats(template, graph, num_inputs):
+    """Compress every other copy of ``template``'s region in ``graph`` with its compiled function."""
+    import rustworkx as rx
+    t_graph = template.get_subgraph()
+    t_output = template.outputs[0]
+    t_anchor = t_graph.predecessors(t_output)[0]
+    port = next(k for k, var in enumerate(t_anchor.outputs) if var is t_output)
+    registered = set(template.recorder.constraints) | set(template.recorder.objectives)
+    anchors = [graph.rxgraph[i] for i in rx.topological_sort(graph.rxgraph)]
+    anchors = [op for op in anchors if type(op) is type(t_anchor) and len(op.outputs) == len(t_anchor.outputs)]
+
+    repeats = []
+    for anchor in anchors:
+        if anchor not in graph.node_table:  # compressed into an earlier copy
+            continue
+        ctx = _MatchContext(
+            template, graph,
+            input_ok=lambda var: True,
+            leaf_ok=lambda var: _is_literal(var) and graph.in_degree(var) == 0,
+            op_ok=lambda op: op in graph.node_table,
+            assume_custom_ops_match=template.assume_custom_ops_match,
+        )
+        state = _walk(_MatchState(), [(t_output, anchor.outputs[port])], ctx)
+        state = _complete(state, ctx) if state is not None else None
+        if state is None or any(var not in state.var_map for var in template.inputs):
+            continue
+
+        # The same checks compress applies to a region: no intermediate used
+        # outside the copy, and no input that depends on the copy itself.
+        ops = set(state.op_map.values())
+        inputs = [state.var_map[var] for var in template.inputs]
+        outputs = [state.var_map[var] for var in template.outputs]
+        internal = {var for op in ops for var in op.outputs}
+        output_set = set(outputs)
+        if internal & set(inputs):
+            continue
+        escapes = any(
+            var in registered or any(user not in ops for user in graph.rxgraph.successors(graph.node_table[var]))
+            for var in internal - output_set
+        )
+        op_indices = {graph.node_table[op] for op in ops}
+        if escapes or any(op_indices & rx.ancestors(graph.rxgraph, graph.node_table[var]) for var in inputs):
+            continue
+        if not _traces_match(state, ctx, inputs):
+            continue
+
+        constants = [c for t, c in state.var_map.items() if t not in set(template.inputs) and not t_graph.predecessors(t)]
+        region = graph.extract_subgraph_nodes(ops | internal | set(inputs) | set(constants))
+        region.name = template.name
+        graph._delete_nodes(
+            [var for var in internal if var not in output_set]
+            + [var for var in constants if graph.out_degree(var) == 0]
+        )
+        copy = JaxCompressedOperation(
+            region, inputs, outputs, name=template.name, device=template.device,
+            jit_kwargs=template.jit_kwargs, compile_separately=template.compile_separately,
+            share_compiled=False, assume_custom_ops_match=template.assume_custom_ops_match)
+        copy.share_compiled = template.share_compiled
+        _share_explicitly(copy, template, list(range(len(inputs))))
+        copy.extra_inputs = inputs[num_inputs:]
+        copy.extra_outputs = outputs[len(outputs) - len(template.extra_outputs):]
+        copy.repeats = []
+        copy.finalize_and_return_outputs(skip_inline=True)
+        repeats.append(copy)
+    return repeats
+
+
 class TestCompressOp(csdl_tests.CSDLTest):
     @staticmethod
     def simple_model(x, y):
@@ -672,6 +1136,40 @@ def step(h, i, p, lib):
                     [csdl_tests.TestingPair(f, f_np, tag=f'{loop_kind} separate={compile_separately}')],
                     verify_derivatives=True)
 
+    def test_repeated_regions_share_compiled_function(self):
+        import csdl_alpha as csdl
+        x_vals = [np.array([0.1, 0.4, 0.7]) + 0.2 * k for k in range(3)]
+        t_vals = [np.array([1.0 + 0.1 * k]) for k in range(3)]
+        c_val = np.array([0.5, -1.0, 2.0])
+        f_np = np.array([sum(np.sum(np.sin(x * t)**2 * c_val) for x, t in zip(x_vals, t_vals))])
+        for compile_separately in [True, False]:
+            self.prep()
+            c = csdl.Variable(name='c', value=c_val)
+            xs = [csdl.Variable(name=f'x{k}', value=v) for k, v in enumerate(x_vals)]
+            ts = [csdl.Variable(name=f't{k}', value=v) for k, v in enumerate(t_vals)]
+            rs = [csdl.sum(csdl.sin(x * t)**2 * c) for x, t in zip(xs, ts)]
+            op = compress([xs[0], ts[0]], rs[0], find_repeats=True, compile_separately=compile_separately)
+            assert len(op.repeats) == 2 and all(r.shared_from is op for r in op.repeats)
+            f = rs[0] + rs[1] + rs[2]
+            self.run_tests(
+                [csdl_tests.TestingPair(f, f_np, tag=f'find_repeats separate={compile_separately}')],
+                verify_derivatives=True)
+
+            # The same with explicit sharing and the inputs in another order.
+            self.prep()
+            c = csdl.Variable(name='c', value=c_val)
+            xs = [csdl.Variable(name=f'x{k}', value=v) for k, v in enumerate(x_vals)]
+            ts = [csdl.Variable(name=f't{k}', value=v) for k, v in enumerate(t_vals)]
+            rs = [csdl.sum(csdl.sin(x * t)**2 * c) for x, t in zip(xs, ts)]
+            op = compress([xs[0], ts[0]], rs[0], share_compiled=False, compile_separately=compile_separately)
+            for k in (1, 2):
+                other = compress([ts[k], xs[k]], rs[k], share_with=op, compile_separately=compile_separately)
+                assert other.shared_from is op
+            f = rs[0] + rs[1] + rs[2]
+            self.run_tests(
+                [csdl_tests.TestingPair(f, f_np, tag=f'share_with separate={compile_separately}')],
+                verify_derivatives=True)
+
 
 if __name__ == '__main__':
     test = TestCompressOp()
diff --git a/csdl_alpha/src/operations/tests/test_compress.py b/csdl_alpha/src/operations/tests/test_compress.py
index 44611ec..488b37b 100644
--- a/csdl_alpha/src/operations/tests/test_compress.py
+++ b/csdl_alpha/src/operations/tests/test_compress.py
@@ -568,3 +568,366 @@ def jax_sim_results(mode):
     body_names = [n.name for n in loops[0].get_subgraph().node_table if isinstance(n, Operation)]
     assert body_names.count('jax_compressed') == 1
     assert not {'sin', 'tanh'} & set(body_names)
+
+
+# ---- Sharing compiled functions between identical regions ----------------------
+
+def _blocks(n, exponent=2.0, scale=1.0, inline=True):
+    """n copies of r = sum(sin(x * t) ** exponent * c) * scale, each with its own x, t; c is shared."""
+    rec = csdl.Recorder(inline=inline)
+    rec.start()
+    c = csdl.Variable(name='c', value=np.array([0.5, -1.0, 2.0]))
+    xs, ts, rs = [], [], []
+    for k in range(n):
+        x = csdl.Variable(name=f'x{k}', value=np.array([0.1, 0.4, 0.7]) + 0.2 * k)
+        t = csdl.Variable(name=f't{k}', value=np.array([1.0 + 0.1 * k]))
+        rs.append(csdl.sum(csdl.sin(x * t) ** exponent * c) * scale)
+        xs.append(x)
+        ts.append(t)
+    return rec, c, xs, ts, rs
+
+
+def _total(rs):
+    return sum(rs[1:], rs[0])
+
+
+def _block_results(n, compress_fn, wrt=lambda xs, ts, c: xs + ts + [c]):
+    rec, c, xs, ts, rs = _blocks(n)
+    ops = compress_fn(xs, ts, rs)
+    f = _total(rs)
+    d = csdl.derivative(f, wrt(xs, ts, c))
+    rec.stop()
+    return ops, rec, [f.value] + [d[w].value for w in wrt(xs, ts, c)]
+
+
+def _assert_all_close(got, expected):
+    for g, e in zip(got, expected):
+        np.testing.assert_allclose(g, e, rtol=1e-10, atol=1e-12)
+
+
+def test_identical_compress_calls_share_one_compiled_function():
+    _, _, expected = _block_results(3, lambda xs, ts, rs: None)
+    ops, rec, got = _block_results(3, lambda xs, ts, rs: [compress([x, t], r) for x, t, r in zip(xs, ts, rs)])
+    _assert_all_close(got, expected)
+    first, *others = ops
+    assert first.shared_from is None
+    assert all(op.shared_from is first for op in others)
+    assert all(op.jit_fn is first.jit_fn for op in others)  # same input order: the very same jitted function
+    rec.execute()                                            # evaluates all three forward operations
+    assert first.jit_fn._cache_size() == 1                   # compiled once for all three
+
+    # Their derivative operations share in the same way.
+    vjps = _vjp_ops(rec)
+    assert len(vjps) == 3 and sum(op.shared_from is None for op in vjps) == 1
+
+
+def _all_nodes(graph):
+    for node in graph.node_table:
+        yield node
+        if isinstance(node, SubgraphOperation) and not isinstance(node, JaxCompressedOperation):
+            yield from _all_nodes(node.get_subgraph())
+
+
+def test_sharing_with_inputs_in_a_different_order():
+    _, _, expected = _block_results(2, lambda xs, ts, rs: None)
+    ops, rec, got = _block_results(2, lambda xs, ts, rs: [compress([xs[0], ts[0]], rs[0]),
+                                                          compress([ts[1], xs[1]], rs[1])])
+    _assert_all_close(got, expected)
+    assert ops[1].shared_from is ops[0] and ops[1]._input_permutation != [0, 1, 2]
+    vjps = _vjp_ops(rec)  # derivative operations are built in the shared order, so they share too
+    assert len(vjps) == 2 and sum(op.shared_from is None for op in vjps) == 1
+
+
+@pytest.mark.parametrize('difference', ['exponent', 'constant'])
+def test_regions_with_different_parameters_do_not_share(difference):
+    rec = csdl.Recorder(inline=True)
+    rec.start()
+    x1 = csdl.Variable(value=np.array([0.3, 0.6]))
+    x2 = csdl.Variable(value=np.array([0.3, 0.6]))
+    if difference == 'exponent':        # a parameter stored on the operation
+        r1, r2 = csdl.sum(csdl.sin(x1) ** 2.0), csdl.sum(csdl.sin(x2) ** 3.0)
+    else:                               # a compiled-in literal constant
+        r1, r2 = csdl.sum(csdl.sin(x1) * 2.0), csdl.sum(csdl.sin(x2) * 3.0)
+    expected = r2.value.copy()
+    op1, op2 = compress(x1, r1, absorb_feeders=True), compress(x2, r2, absorb_feeders=True)
+    rec.stop()
+    assert op2.shared_from is None
+    sim = csdl.experimental.PySimulator(rec)
+    sim.run()
+    np.testing.assert_allclose(sim[r2], expected)
+
+
+class _Scale(csdl.CustomExplicitOperation):
+    def __init__(self, k):
+        super().__init__()
+        self.k = k
+
+    def evaluate(self, x):
+        self.declare_input('x', x)
+        f = self.create_output('f', x.shape)
+        self.declare_derivative_parameters('f', 'x')
+        return f
+
+    def compute(self, input_vals, output_vals):
+        output_vals['f'] = self.k * input_vals['x']**2
+
+    def compute_derivatives(self, input_vals, output_vals, derivatives):
+        derivatives['f', 'x'] = np.diag(2 * self.k * input_vals['x'])
+
+
+@pytest.mark.parametrize('assume', [False, True])
+def test_custom_operations_share_only_when_assumed(assume):
+    ks = [2.0, 2.0, 5.0]
+    def build(do_compress):
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        xs = [csdl.Variable(value=np.array([0.2, 0.9])) for _ in ks]
+        rs = [csdl.sum(csdl.sin(_Scale(k).evaluate(x))) for k, x in zip(ks, xs)]
+        # Assume only for the two that really match; the k=5 one never assumes.
+        ops = [compress(x, r, assume_custom_ops_match=(assume and k == 2.0)) for k, x, r in zip(ks, xs, rs)] if do_compress else None
+        f = _total(rs)
+        g = csdl.derivative(f, xs)
+        rec.stop()
+        return ops, [f.value] + [g[x].value for x in xs]
+
+    _, expected = build(False)
+    ops, got = build(True)
+    _assert_all_close(got, expected)
+    assert (ops[1].shared_from is ops[0]) == assume
+    assert ops[2].shared_from is None
+
+
+def test_find_repeats_compresses_every_copy():
+    _, _, expected = _block_results(3, lambda xs, ts, rs: None)
+    ops, rec, got = _block_results(3, lambda xs, ts, rs: [compress([xs[0], ts[0]], rs[0], find_repeats=True)])
+    _assert_all_close(got, expected)
+    template = ops[0]
+    assert len(template.repeats) == 2
+    assert all(op.shared_from is template for op in template.repeats)
+    names = [n.name for n in rec.active_graph.node_table if hasattr(n, 'inputs')]
+    assert names.count('jax_compressed') == 3 and 'sin' not in names
+
+
+def test_find_repeats_in_a_chain_takes_non_overlapping_copies():
+    def build(do_compress, layers_per_copy):
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        x = csdl.Variable(value=np.array([0.2, 0.5]))
+        hs, h = [], x
+        for _ in range(6):
+            h = csdl.tanh(h * 0.9) + 0.05
+            hs.append(h)
+        op = compress(x, hs[layers_per_copy - 1], find_repeats=True) if do_compress else None
+        f = csdl.sum(h**2)
+        g = csdl.derivative(f, x)
+        rec.stop()
+        return op, [f.value, g.value]
+
+    _, expected = build(False, 1)
+    for layers_per_copy, num_repeats in [(1, 5), (2, 2), (4, 0)]:
+        op, got = build(True, layers_per_copy)
+        _assert_all_close(got, expected)
+        assert len(op.repeats) == num_repeats
+
+
+def test_find_repeats_skips_a_copy_with_an_escaping_intermediate():
+    rec, c, xs, ts, rs = _blocks(3)
+    # In block 1, the intermediate x1 * t1 is also used outside the block.
+    leak = [n for n in rec.active_graph.node_table if hasattr(n, 'inputs') and xs[1] in n.inputs][0].outputs[0]
+    w = csdl.sum(leak)
+    op = compress([xs[0], ts[0]], rs[0], find_repeats=True)
+    rec.stop()
+    assert len(op.repeats) == 1 and op.repeats[0].outputs == [rs[2]]
+
+
+def test_find_repeats_with_outputs_on_independent_paths():
+    # Output 1 is not upstream of output 0, so its operations are found by
+    # walking forward from the shared input.
+    def build(do_compress):
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        xs = [csdl.Variable(value=np.array([0.3, 0.8]) + k) for k in range(3)]
+        outs = [(csdl.sum(csdl.sin(x)), csdl.cos(x) * 2.0) for x in xs]
+        op = compress(xs[0], list(outs[0]), find_repeats=True) if do_compress else None
+        f = sum((r + csdl.sum(q**2) for r, q in outs[1:]), outs[0][0] + csdl.sum(outs[0][1]**2))
+        g = csdl.derivative(f, xs)
+        rec.stop()
+        return op, [f.value] + [g[x].value for x in xs]
+
+    _, expected = build(False)
+    op, got = build(True)
+    _assert_all_close(got, expected)
+    assert len(op.repeats) == 2
+
+
+@pytest.mark.parametrize('compile_separately', [True, False])
+def test_find_repeats_under_jax_simulator(compile_separately):
+    def results(do_compress):
+        rec, c, xs, ts, rs = _blocks(3, inline=False)
+        op = compress([xs[0], ts[0]], rs[0], find_repeats=True, compile_separately=compile_separately) if do_compress else None
+        f = _total(rs)
+        rec.stop()
+        sim = csdl.experimental.JaxSimulator(rec, gpu=False, additional_inputs=xs + ts + [c], additional_outputs=[f])
+        sim.run()
+        d = sim.compute_totals()
+        return op, [sim[f]] + [d[f, w] for w in xs + ts + [c]]
+
+    _, expected = results(False)
+    op, got = results(True)
+    _assert_all_close(got, expected)
+    assert len(op.repeats) == 2
+    if compile_separately:  # all three copies called one compiled function
+        assert op.jit_fn._cache_size() == 1
+
+
+# ---- Explicit sharing: compress(..., share_with=op) ------------------------------
+
+def _vjp_ops(rec):
+    # An operation can sit in more than one graph (a derivative loop's body reuses it), so dedupe.
+    nodes = dict.fromkeys(_all_nodes(rec.active_graph))
+    return [n for n in nodes if isinstance(n, JaxCompressedOperation) and n.name.startswith('vjp')]
+
+
+def test_share_with_reuses_function_and_derivatives():
+    _, _, expected = _block_results(2, lambda xs, ts, rs: None)
+    def compress_both(xs, ts, rs):
+        first = compress([xs[0], ts[0]], rs[0], share_compiled=False)
+        second = compress([ts[1], xs[1]], rs[1], share_compiled=False, share_with=first)
+        return [first, second]
+    ops, rec, got = _block_results(2, compress_both)
+    _assert_all_close(got, expected)
+    assert ops[1].shared_from is ops[0]
+    vjps = _vjp_ops(rec)
+    assert len(vjps) == 2 and sum(op.shared_from is None for op in vjps) == 1
+
+
+def test_share_with_second_derivatives_share_too():
+    def build(do_compress):
+        rec, c, xs, ts, rs = _blocks(2)
+        if do_compress:
+            first = compress([xs[0], ts[0]], rs[0], share_compiled=False)
+            compress([xs[1], ts[1]], rs[1], share_compiled=False, share_with=first)
+        f = _total(rs)
+        g = csdl.derivative(f, xs)
+        h = csdl.derivative(g[xs[0]] + g[xs[1]], xs)
+        rec.stop()
+        return rec, [h[x].value for x in xs]
+
+    _, expected = build(False)
+    rec, got = build(True)
+    _assert_all_close(got, expected)
+    vjps = _vjp_ops(rec)
+    first_order = [op for op in vjps if not op.name.startswith('vjp_vjp')]
+    second_order = [op for op in vjps if op.name.startswith('vjp_vjp')]
+    assert len(first_order) == 2 and sum(op.shared_from is None for op in first_order) == 1
+    assert len(second_order) == 2 and sum(op.shared_from is None for op in second_order) == 1
+
+
+def test_share_with_custom_operations():
+    def build(do_compress):
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        xs = [csdl.Variable(value=np.array([0.2, 0.9]) + k) for k in range(2)]
+        rs = [csdl.sum(csdl.sin(_Scale(2.0).evaluate(x))) for x in xs]
+        ops = None
+        if do_compress:
+            first = compress(xs[0], rs[0])
+            ops = [first, compress(xs[1], rs[1], share_with=first)]  # no assume_custom_ops_match needed
+        f = _total(rs)
+        g = csdl.derivative(f, xs)
+        rec.stop()
+        return ops, rec, [f.value] + [g[x].value for x in xs]
+
+    _, _, expected = build(False)
+    ops, rec, got = build(True)
+    _assert_all_close(got, expected)
+    assert ops[1].shared_from is ops[0]
+    vjps = _vjp_ops(rec)
+    assert len(vjps) == 2 and sum(op.shared_from is None for op in vjps) == 1
+
+
+def test_share_with_mismatch_raises_before_changing_graph():
+    rec = csdl.Recorder(inline=True)
+    rec.start()
+    x1 = csdl.Variable(value=np.array([0.3, 0.6]))
+    x2 = csdl.Variable(value=np.array([0.3, 0.6]))
+    r1, r2 = csdl.sum(csdl.sin(x1) ** 2.0), csdl.sum(csdl.sin(x2) ** 3.0)
+    first = compress(x1, r1)
+    before = op_names(rec)
+    with pytest.raises(ValueError, match='does not match the region of share_with'):
+        compress(x2, r2, share_with=first)
+    with pytest.raises(TypeError, match='share_with must be'):
+        compress(x2, r2, share_with='first')
+    assert op_names(rec) == before
+    rec.stop()
+
+
+def test_share_with_takes_jit_kwargs_from_the_shared_operation():
+    rec, c, xs, ts, rs = _blocks(2)
+    first = compress([xs[0], ts[0]], rs[0], jit_kwargs={'keep_unused': True})
+    with pytest.raises(ValueError, match='jit_kwargs differ'):
+        compress([xs[1], ts[1]], rs[1], share_with=first, jit_kwargs={'keep_unused': False})
+    second = compress([xs[1], ts[1]], rs[1], share_with=first)
+    rec.stop()
+    assert second.jit_kwargs == {'keep_unused': True} and second.shared_from is first
+
+
+@pytest.mark.parametrize('how', ['automatic', 'share_with'])
+def test_derivatives_share_when_blocks_are_compressed_as_they_are_recorded(how):
+    # Compressing a block frees node indices that later blocks reuse, so later
+    # blocks are numbered differently from the first. Sharing renumbers each
+    # matched subgraph like its template, so CSDL records their derivative
+    # graphs identically and those share one compiled function too.
+    def build(do_compress):
+        rec = csdl.Recorder(inline=True)
+        rec.start()
+        x = csdl.Variable(value=np.array([0.2, 0.5, 0.9]))
+        h, first = x, None
+        for _ in range(4):
+            start = h
+            for _ in range(5):
+                h = csdl.tanh(0.9 * h + 0.1) + 0.05 * csdl.sin(h)
+            if do_compress:
+                op = compress(start, h, absorb_feeders=True, share_with=first if how == 'share_with' else None)
+                first = first or op
+        f = csdl.sum(h**2)
+        g = csdl.derivative(f, x)
+        rec.stop()
+        return rec, [f.value, g.value]
+
+    _, expected = build(False)
+    rec, got = build(True)
+    _assert_all_close(got, expected)
+    vjps = _vjp_ops(rec)
+    owners = {id(op.shared_from or op) for op in vjps}
+    assert len(vjps) == 4 and len(owners) == 1
+    rec.execute()
+    owner = vjps[0].shared_from or vjps[0]
+    assert owner.jit_fn._cache_size() == 1
+
+
+def test_derivatives_of_shared_operations_share_without_matching(monkeypatch):
+    import csdl_alpha.src.operations.compress_operations as compress_module
+    rec, c, xs, ts, rs = _blocks(3)
+    first = compress([xs[0], ts[0]], rs[0], share_compiled=False)
+    for k in (1, 2):
+        compress([ts[k], xs[k]], rs[k], share_with=first)
+
+    calls = {'match': 0, 'trace': 0}
+    real_match, real_trace = compress_module._match, compress_module._trace_region
+    def counting_match(*args, **kwargs):
+        calls['match'] += 1
+        return real_match(*args, **kwargs)
+    def counting_trace(*args, **kwargs):
+        calls['trace'] += 1
+        return real_trace(*args, **kwargs)
+    monkeypatch.setattr(compress_module, '_match', counting_match)
+    monkeypatch.setattr(compress_module, '_trace_region', counting_trace)
+
+    f = _total(rs)
+    csdl.derivative(f, xs + ts)
+    csdl.derivative(f, xs + ts)  # a second derivative call reuses the first's derivative function
+    rec.stop()
+    assert calls == {'match': 0, 'trace': 0}
+    vjps = _vjp_ops(rec)
+    assert len(vjps) == 6 and len({id(op.shared_from or op) for op in vjps}) == 1
diff --git a/csdl_alpha/src/recorder.py b/csdl_alpha/src/recorder.py
index 0c98817..5991c76 100644
--- a/csdl_alpha/src/recorder.py
+++ b/csdl_alpha/src/recorder.py
@@ -58,6 +58,11 @@ def __init__(self,
         self.constraints:dict[Variable, tuple[np.array, np.array, np.array]] = {}
         self.objectives:dict[Variable, np.array] = {}
 
+        # Compressed operations that own a compiled function, by a cheap
+        # structural key, so identical regions can share it (see
+        # csdl_alpha.src.operations.compress_operations.compress).
+        self.compressed_operations:dict[tuple, list] = {}
+
         self.namespace_tree = Namespace(None)
         self.active_namespace = self.namespace_tree
 

From 0f99cb83a447cd8f529bb33bd8200d8d5412a7a1 Mon Sep 17 00:00:00 2001
From: Michael Warner <hgxe14@gmail.com>
Date: Thu, 8 Oct 2026 18:59:03 -0700
Subject: [PATCH 4/4] Ignore memory addresses when comparing region traces

JAX 0.4 and 0.6 print a pure_callback's Python function with its memory
address, so regions with custom operations never matched under
share_with or assume_custom_ops_match on the CI's JAX versions. Whether a
callback may match is decided separately, so the address carries nothing.
---
 csdl_alpha/src/operations/compress_operations.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/csdl_alpha/src/operations/compress_operations.py b/csdl_alpha/src/operations/compress_operations.py
index 37c3700..73fa4b6 100644
--- a/csdl_alpha/src/operations/compress_operations.py
+++ b/csdl_alpha/src/operations/compress_operations.py
@@ -3,6 +3,7 @@
 from csdl_alpha.src.operations.operation_subclasses import SubgraphOperation
 import csdl_alpha.utils.testing_utils as csdl_tests
 import numpy as np
+import re
 
 class CompressedOperation(SubgraphOperation):
     def __init__(self, subgraph, inputs, outputs, name, jax_jit=True):
@@ -818,7 +819,11 @@ def region(*args):
         closed = jax.make_jaxpr(region)(*(jax.ShapeDtypeStruct(var.shape, np.float64) for var in inputs))
     except Exception:
         return None
-    return str(closed.jaxpr), tuple((np.shape(c), np.asarray(c).tobytes()) for c in closed.consts)
+    # Some JAX versions print a callback's Python function with its memory
+    # address, which differs between otherwise identical regions. Whether a
+    # callback may match at all is decided separately (assume_custom_ops_match).
+    text = re.sub(r'0x[0-9a-fA-F]+', '0x?', str(closed.jaxpr))
+    return text, tuple((np.shape(c), np.asarray(c).tobytes()) for c in closed.consts)
 
 
 def _traces_match(state, ctx, c_inputs):