Add csdl.experimental.compress for JAX-compiled subgraph operations

[HgXe]

Human TLDR: Automatically wrapping sections of a CSDL model in a jax function, like we previously discussed with a custom operation, except this uses the CompressedOperation class instead. By default it will put that inside a jax pure_callback, meaning it's compiled completely separately. Seems to work pretty well, and save a decent amount of compile time in the toy cases I've tested on.

Adds csdl.experimental.compress(inputs, outputs, ...). It replaces the operations between inputs and outputs with one JaxCompressedOperation that runs as a single jitted JAX function and is still differentiable to any order.

z = model(x, y)                                   # any CSDL code
op = csdl.experimental.compress([x, y], z)        # the region becomes one operation
f = 2.0 * z                                       # downstream code keeps using z
dfdx = csdl.derivative(f, x)

The main motivation is long JAX compile times and compile memory on large models. By default each compressed region compiles on its own, outside JaxSimulator's program, and can run on its own device.

What it does

compress finds the operations between inputs and outputs with Graph._get_intersection, moves them into a subgraph with Graph.extract_subgraph_nodes, and attaches one operation that produces the same output Variable objects. Downstream operations, constraints, objectives, and simulator handles keep working. Any other variable the region reads, such as a parameter, becomes an extra input read from the graph, so nothing is frozen. Only literal constants (csdl.Constant) are compiled in.

JaxCompressedOperation subclasses CompressedOperation. CompressedOperation itself is unchanged, so its inline path still does not need JAX. The subclass differs in three ways:

	CompressedOperation	JaxCompressedOperation
Inline / PySimulator	Executes the subgraph op by op	Calls one jax.jit function
Reverse mode	Records CSDL's VJP ops into the outer graph	Records CSDL's VJP into a copy of the subgraph and wraps it as another JaxCompressedOperation: one compiled node, itself differentiable
Under JaxSimulator	Traced into the simulator's program	By default a host callback to its own compiled function (compile_separately=True); False traces it in

Derivatives use CSDL's own reverse mode, not jax.vjp. So custom operations use their compute_derivatives (JAX cannot differentiate their host callbacks), implicit operations use their solver's implicit derivative (JAX cannot reverse-differentiate the solver's while_loop), and Hessians work. The derivative operation inherits the forward operation's options.

Options

Option	Default	Effect
absorb_feeders	False	Also compress upstream operations whose outputs only the region uses, such as the broadcast of a scalar constant
intermediates	'none'	'none': only outputs, and an intermediate used outside the region raises. 'used': add those intermediates as outputs. 'all': every variable the region computes is an output
device	JAX default	'cpu', 'gpu', or a jax.Device
jit_kwargs	{}	Passed to jax.jit
compile_separately	True	Under JaxSimulator, call the operation through a host callback (backends.jax.utils.host_callback) instead of compiling it into the simulator's program
share_compiled	True	Reuse the compiled function of an earlier identical compressed region in this recorder
find_repeats	False	Also compress every other copy of the region in the graph with the same compiled function (op.repeats)
share_with	None	Reuse a given compressed operation's compiled function; the region is checked and a mismatch raises before the graph changes
assume_custom_ops_match	False	Let regions containing custom operations (or other callbacks) count as identical

Why separate compilation is the default

With compile_separately=True, the simulator's program contains only a callback, so its XLA compile excludes the region, and the region runs on its own device. Each call costs a round trip through host memory and loses XLA optimization across the boundary. Under vmap (csdl.derivative(..., loop=False)), the callback runs once per batch element.

On a chain of tanh/sin layers split into 100-layer compressed blocks (CPU, JaxSimulator; setup plus first run and compute_totals, then steady state):

Layers	Uncompressed	Fused blocks	Separate blocks
400	4.6 s, then 1.9 ms	4.3 s, then 1.9 ms	3.3 s, then 3.2 ms
1200	88 s, then 965 ms	87 s, then 947 ms	9.9 s, then 8.9 ms

At 1200 layers the single fused program also hits an XLA-CPU slowdown on very long unrolled chains, which the separate blocks stay under, so the steady-state gap there is larger than usual. On a GPU JaxSimulator, two separate blocks placed on 'gpu' and 'cpu' each ran on their own device and matched the uncompressed model.

Note - Fused blocks means compile_seperately=False, so we essentially end up with the same Jax function as uncompressed (hence the numbers are essentially the same).

Sharing compiled functions between identical regions

Identical regions (repeated components, blocks of the same layers) can share one compiled function, so with compile_separately each distinct region compiles once however many copies there are.

• When regions are identical. A structural walk pairs the two regions' operations port by port: same class and shapes, each input wired from the paired producer's output at the same port, and equal compiled-in constants. Each variable has one producer and inputs are ordered, so the walk goes back from the outputs with no search. Then each region is traced once with jax.make_jaxpr, running the paired operations in the same order, and the jaxprs must be equal. That catches parameters the walk cannot see, such as reshape shapes or exponents. Inputs may be in a different order; the operation reorders its arguments to the shared function.
• Custom operations. A custom operation is a host callback, and JAX cannot see the Python state it computes with (two instances with different stored parameters give identical jaxprs). So regions containing a callback only match with assume_custom_ops_match=True, or for the pair named by share_with, which asserts it.
• Three ways to share. Automatic (share_compiled, checked against earlier compressed operations in Recorder.compressed_operations), explicit (share_with=op), and search (find_repeats, which finds the other copies in the graph). With find_repeats, overlapping copies (a two-sin region in a chain of sins) are taken greedily in topological order, which gives the most copies when overlaps line up like a chain; share_with gives exact control.
• Derivatives. Operations that share a function form a group. A derivative operation is shared within its group by which outputs are seeded and which inputs are differentiated, with no further matching, since it differentiates the same function. Each member still records its own CSDL derivative subgraph, which higher derivatives need.

Same chain benchmark, now as identical 100-layer blocks with compile_separately=True (CPU, JaxSimulator; graph building, first run, and first compute_totals, which include compiling; then steady state):

Blocks	No sharing	Shared	find_repeats
4	3.6 s, then 3.6 ms	1.2 s, then 3.1 ms	1.2 s, then 3.2 ms
12	10.3 s, then 8.7 ms	1.7 s, then 8.1 ms	2.0 s, then 8.3 ms
24	20.1 s, then 16.8 ms	2.7 s, then 16.4 ms	3.1 s, then 16.1 ms

With sharing, one forward and one derivative function compile regardless of the block count; what remains is CSDL graph and derivative recording plus about 0.04 s of matching per block.

Tests

• TestJaxCompress in compress_operations.py uses CSDLTest: values against NumPy references, and finite-difference derivatives, for each option, intermediates, a Newton implicit solve, a custom operation, a loop, nesting, compressing inside frange and enter_loop bodies, and the docstring example. It passes in all eight combinations of run_test_matrix.sh.
• src/operations/tests/test_compress.py compares compressed and uncompressed models exactly, for gradients, Hessians, both simulators in both compile modes, the option validations and errors, device and jit_kwargs inheritance, the callback replacing the region in the traced program, and compressing inside frange and enter_loop bodies (with and without the loop index, both compile modes).
• Sharing tests (in test_compress.py, plus a CSDLTest case for find_repeats and share_with): separate compress calls, reordered inputs, share_with (including mismatches, custom operations, second derivatives), regions differing in a parameter or constant not sharing, custom operations sharing only when asserted, find_repeats (copies, overlapping chains, escaping intermediates, outputs on independent paths, JaxSimulator in both compile modes), derivative operations sharing without any matching or tracing, and each shared function compiling once.
• Full suite: 375 passed, 2 failed. Both failures also occur on main with NumPy 2 and are unrelated (test_variable.py::test_docstrings expects True and gets np.True_; test_custom_implicit.py::test_custom_implicit raises "setting an array element with a sequence").

Notes

• compress uses private graph and recorder methods: Graph._get_intersection, Graph._delete_nodes, and Recorder._enter_subgraph.
• Recorder gains one attribute, compressed_operations, the per-recorder registry used for automatic sharing.
• Each csdl.derivative call records the derivative into a copy of the region's graph. Operations in a sharing group, and repeated derivative calls with the same seeded outputs and differentiated inputs, reuse one compiled derivative function. That function is the first member's, so a member's derivative values can differ from its own recorded graph's in the last bits (summation order). The derivative recomputes the forward values it needs; when fused, XLA removes the duplicate work.
• Random operations are not supported inside a region. Calling compress inside a loop body (csdl.frange or enter_loop) is tested; inside a composed operation it is untested.

[msperryucsd]

Looks good to me. As we discussed, we should consider adding the CSDL-native hashing mechanism to remove the jax_pr dependency.