Skip to content

Control-Flow Body Inputs (ONNX)

Our IR-only pipeline still emits ONNX control-flow nodes, so plugin authors and optimizer work must respect the schema-defined subgraph contracts. Use this guide when touching converter/ir_context.py, the control-flow plugins, or the structural tests that verify body graphs (for example tests/extra_tests/scan/).

TL;DR

  • If exposes exactly one explicit input (cond). Branch graphs declare zero formal inputs and capture any extra tensors by name from the parent scope.
  • Loop forwards everything explicitly: iteration_num, condition, and N loop-carried values are provided to the body each iteration. Thread outer tensors through the carried tuple when they are not constants.
  • jax.lax.scan is currently lowered through ONNX Loop, not the ONNX Scan operator. Scanned inputs are passed explicitly into the loop body and sliced by iteration; nothing is implicitly captured.

If: implicit capture only

The ONNX If node accepts a single boolean input. Both then_branch and else_branch Graph attributes must list zero inputs, so any tensor needed inside the branch is referenced directly by name and resolved from the enclosing graph. Exporters that try to add branch parameters fail schema validation (input size of if-op should be 1). Our lowering honours that contract—branch builders leave the input list empty and rely on IRContext to resolve captured values.

Child contexts use create_child_context(scope="control_flow", ...). Only If opts into implicit_captures=True. require_value_for_var first checks the local binding, then resolves the same JAX variable in the enclosing context. The captured value is mirrored by an Identity into a branch-local value and cached by variable identity. This preserves computed flags and tensor keyword arguments passed to nested ONNX Functions, regardless of the keyword name. An argument such as deterministic=jnp.logical_not(deterministic) must retain the negation; its name cannot substitute for its value.

At the JAX tracing boundary, dynamic Function keywords are regular primitive operands. Their dependencies participate in branch inputs and dead-code elimination before ONNX lowering. Storing a tracer only in primitive parameters would let JAX discard its producer, leaving no variable binding to capture.

Declared external BOOL parameters can also be resolved lazily by name when no JAX variable is supplied. Each nested Function receives captures through fresh formal inputs. No parent parameter-to-IR-value mapping is copied into a child. The initial trace retains named input_params promotion; Function-body tracing disables promotion of explicit Python values because forwarded flags already have operands. Control-flow contexts inherit that policy. Explicit True and False inside a Function therefore remain constants across nested If scopes. Omitted keywords in a Function body use the callee's default values instead of implicitly capturing a same-named export parameter. Function and Loop/scan contexts keep implicit captures disabled and require explicit inputs; an unknown variable raises an error rather than creating an unbound placeholder or selecting a similarly named parameter.

Loop: explicit interface

Loop nodes receive an optional trip-count M, an optional initial condition, and N loop-carried initial values. Consequently, the body graph declares iteration_num, condition, and those N carried values as explicit inputs. Each iteration returns a continuation flag, the updated carried values, and optional scan outputs that map back to the parent node. When the body needs an outer tensor that is not constant, add it to the carried tuple (passing it through unchanged if necessary). This keeps dependencies explicit and aligns with ONNX Runtime’s validation. Our fori_loop and scan plugins rely on this ordering when constructing body graphs.

LoopCallParameters handles declared export flags that were Python values during tracing and therefore have no JAX operand. The first actual request adds a parent flag input, a fresh Loop-body formal, and an unchanged carried output. These internal outputs precede scan outputs and are removed before binding the JAX results. Nested Ifs and Loops can request the same flag through their enclosing context; unused parameters add no model inputs or Loop state. This mechanism never enables general implicit captures for Loop bodies.

fori_loop also carries closed-over JAX values as regular primitive operands and unchanged Loop state. Traced arrays are not stored as static primitive parameters or converted to NumPy during body construction.

jax.lax.scan: Loop lowering with explicit sequence state

jax.lax.scan lowers to ONNX Loop. When xs=None, the loop body receives the standard Loop header plus explicit constant/carry state. When scan inputs are present, the outer sequence tensors are passed explicitly as loop-carried state; the body gathers the current per-step slice by iteration index and returns the sequence state unchanged. It yields updated carry values plus scan outputs that the parent graph stacks or reshapes as required. Tests under tests/extra_tests/scan/ assert these invariants so regressions surface quickly.

Practical tips

  • Bind carried/captured operands in the child explicitly. Read them with require_value_for_var, which also materializes literals in the current scope. get_value_for_var remains a compatibility convenience.
  • When authoring tests, assert Loop node/body arity directly for control-flow body contracts; use expect_graph for top-level structural paths and imported ONNX Function bodies.
  • If ONNX adjusts the schemas, update the converter code and this guide together so plugin authors retain a single source of truth.

Constants inside subgraphs

  • ONNX Functions cannot contain graph initializers, but control-flow attributes are GraphProto values and may legally contain them. During construction, constants inside Loop bodies and Function graphs are emitted as Constant nodes; optimization may later lift control-flow constants into subgraph initializers.
  • Our converter runs subgraph construction in “function mode”, which makes builder initializer helpers produce Constant nodes instead of registering initializers immediately. Imported ONNX Functions remain initializer-free.
  • Plugin authors should always use ctx.builder.add_initializer_from_* or ctx.bind_const_for_var(...) for constants so the correct form is emitted automatically in subgraphs. Avoid writing to any _initializers lists directly.
  • Post‑processing loosens shapes inside subgraphs: value shapes in Loop bodies are set to rank‑only (all dims unknown) to reduce schema friction and improve portability. Structural tests assert this behaviour.