Sequences, maps and subgraphs#

The algorithm overview describes how onnx_optim infers the element type and shape of every tensor value in a graph. ONNX also carries sequence and map values, and nests graphs inside If / Loop / Scan nodes and model-local functions. This page describes how those richer values and nested graphs are handled.

Sequence values#

A tensor descriptor (OptimTensor) cannot represent the output of SequenceConstruct, SequenceEmpty, SplitToSequence and the other sequence operators, because those values are sequences of tensors rather than tensors. They are described by a separate OptimSequence descriptor (see onnx_light/onnx_optim/optim_sequence.h), the sequence analogue of OptimTensor. It records:

  • the common element dtype shared by every tensor of the sequence (kUndefined when unknown, e.g. for an empty SequenceConstruct);

  • either one OptimShape per element (when the per-element shapes are known) or a single, possibly symbolic, length OptimDim (e.g. for SplitToSequence whose length depends on a runtime split value).

OptimSequence::HasElemDtype() and OptimSequence::HasElemShapes() distinguish unknown from empty sequence (length 0).

ShapesContext therefore keeps two parallel maps: the name → OptimTensor map for tensors and a name → OptimSequence map for sequences. The sequence map has its own C++ accessors (SetSequence / HasSequence / GetSequence / ShapesContext::Sequences()); the Python bindings expose the read-only introspection helpers has_sequence / sequence_names / sequences_size. CheckInputsAvailable and CheckOutputsNotAvailable consult both maps, so a value is considered “known” whether it is a tensor or a sequence.

Map values#

ONNX map values (map(K, V)) and sequences of maps (seq(map(K, V))) are produced by the traditional-ML operators (ZipMap, CastMap, DictVectorizer, …). They are not given a dedicated descriptor; instead they reuse OptimTensor with a map-valued TensorType enumerator (see onnx_light/onnx_op/light_op_schema.h):

  • kMapStringInt64 / kMapInt64String / kMapInt64Float / … — the map(K, V) element types;

  • kSeqMapStringFloat / kSeqMapInt64Float — the seq(map(K, V)) types, for example the output of ZipMap, whose leading (sequence-length) dimension is the input batch size.

Because the element type already encodes the key/value pair, map-valued descriptors flow through the same tensor map and the same write-back path as ordinary tensors; only the dtype differs.

Subgraphs#

Control-flow operators embed graphs in their attributes: then_branch / else_branch for If, body for Loop and Scan. Each subgraph is inferred in a child ShapesContext derived from the parent (see onnx_light/onnx_optim/shapes/controlflow/shape_controlflow.cc):

  1. The child context is copied from the parent so the subgraph can read outer-scope values it captures by name, then seeded with the subgraph’s own initializers and formal inputs (the loop iteration count / condition, the scan slices, …).

  2. ShapesContext::set_current_subgraph() records the owning node index and the attribute name ("body", "then_branch" or "else_branch") so that every ShapeEvent recorded inside carries its subgraph_node_index / subgraph_attr_name and can be traced back to its originating subgraph (see Shape-inference events (ShapeEvent)).

  3. ComputeShapes runs over the subgraph nodes in the child context.

  4. The subgraph outputs are read back from the child context and mapped onto the control-flow node’s outputs. If merges the matching then / else outputs (a differing dimension becomes a fresh symbol bounded above by the max of the two branch expressions); Loop / Scan reshape the body outputs into the loop-carried and scan-output shapes.

  5. The child context is retained on the parent through ShapesContext::RegisterSubgraphContext(), keyed by the owning node index and the attribute name, so the descriptors inferred inside the subgraph stay inspectable once the control-flow node has been processed. They are reachable through ShapesContext::HasSubgraphContext(), ShapesContext::GetSubgraphContext() and ShapesContext::SubgraphContexts() (exposed in Python as has_subgraph_context / subgraph_context / subgraph_context_keys / subgraph_contexts_size).

Local functions#

A call to a model-local function (a FunctionProto listed in model.functions) is expanded by ExpandLocalFunctionCall (in onnx_light/onnx_optim/shapes/shape_inference.cc). Unlike a subgraph, a function body does not capture outer-scope values, so it runs in a fresh, isolated child ShapesContext that:

  • inherits the caller’s opset versions, then lets the function’s own opset_import override them;

  • receives a copy of the parent’s local-function map, so nested local calls are dispatched too;

  • binds the call-site inputs to the function’s formal input names positionally (carrying either an OptimTensor or an OptimSequence descriptor);

  • resolves ref_attr_name attribute references in the body against the call-site node’s attributes before inference.

After the body is inferred, the function outputs are copied back onto the caller-visible output names.

Note

Each If / Loop / Scan node retains the child ShapesContext it builds for every attribute subgraph (then_branch / else_branch / body) on its parent context, keyed by (node_index, attr_name), so the inferred descriptors inside the subgraph remain inspectable (and traceable through the event log) after inference completes. Local-function calls still build a child context that is discarded once its outputs have been merged back; extending the same retention to ExpandLocalFunctionCall is a natural follow-up.

API reference#