onnx_light.onnx_optim.shape_inference#

Top-level shape-inference helpers backed by a C++ library.

This module exposes the onnx_optim shape-inference pipeline at the ModelProto level. It mirrors the C++ helpers ComputeShapeModel() and ApplyInferredShapesToModel() combined into a single infer_shapes_model() entry point.

Typical usage:

from onnx_light.onnx_optim.shape_inference import infer_shapes_model

infer_shapes_model(model)
# model.graph.output[i].type and model.graph.value_info now carry
# the inferred element types and shapes.

The module is exposed as onnx_light.onnx_optim.shape_inference.

class onnx_light.onnx_optim.shape_inference.ComputeContext(*args, **kwargs)#

Holds the in-place reuse opportunities computed for a graph, mirroring the way ShapesContext holds inferred descriptors. Populate it with compute_inplace_reuse_graph (consuming a ShapesContext), then read the result through reuse / node_reuse / memory or persist it with write_to_metadata.

clear(self) → None#: Empties the stored result.

clear_events(self) → None#: Empties the decision log without touching computed reuse results.

compute_inplace_reuse_graph(self, graph: onnx_light.onnx_py._onnxpyprotoop.GraphProto, ctx: onnx_light.onnx_py._onnxpyoptim.shape_inference.ShapesContext, allow_input_overwrite: bool = False, value_tags: collections.abc.Mapping[str, str] = {}, verbose: int = 0) → None#

Guesses, for every node of graph, which outputs reuse which input buffers in place, using the shapes already inferred into ctx, and stores the result in this context (replacing any previous result).

By default declared graph inputs are never overwritten in place; set allow_input_overwrite=True to let an input be reused like an intermediate.

When value_tags is provided (a {name: tag} dict such as the one returned by compute_value_and_node_tags()), released values that carry the "shape" tag are also stored in release_after_shape_tagged and written to onnx_light.release_after_shape_tag by write_to_metadata(). When value_tags is omitted, this method reuses the last tags stored by compute_value_and_node_tags() on the same context, if any. verbose is currently accepted for API compatibility and has no effect.

compute_release_after_shape_tagged(self, verbose: int = 0) → list[list[str]]#: Returns release_after_shape_tagged; the optional verbose argument is accepted for API compatibility.

compute_value_and_node_tags(self, graph: onnx_light.onnx_py._onnxpyprotoop.GraphProto, verbose: int = 0) → tuple#

compute_value_and_node_tags(self, function: onnx_light.onnx_py._onnxpyprotoop.FunctionProto, verbose: int = 0) → tuple

compute_value_and_node_tags(self, nodes: list, verbose: int = 0) → tuple

Overloaded function.

compute_value_and_node_tags(self, graph: onnx_light.onnx_py._onnxpyprotoop.GraphProto, verbose: int = 0) -> tuple

Computes semantic shape/axes/weight/ambiguous tags for values and nodes in graph and stores the result in this context. verbose is currently accepted for API compatibility and has no effect.

compute_value_and_node_tags(self, function: onnx_light.onnx_py._onnxpyprotoop.FunctionProto, verbose: int = 0) -> tuple

Computes semantic shape/axes/weight/ambiguous tags for values and nodes in function and stores the result in this context. verbose is currently accepted for API compatibility and has no effect.

compute_value_and_node_tags(self, nodes: list, verbose: int = 0) -> tuple

Computes semantic shape/axes/weight/ambiguous tags for a node list and stores the result in this context. verbose is currently accepted for API compatibility and has no effect.

events(self) → list#: Returns the append-only decision log as a list of ComputeEvent entries.

property events_enabled#: When True, compute_inplace_reuse_graph() appends one ComputeEvent per in-place decision, release decision and shape-tagged release decision. Default is False.

property memory#: The per-node memory snapshots as a list (one entry per node, same order as graph.node). Each entry is a dict[str, int | str | dict[str, int | str]] describing the live input/initializer/intermediate buffers plus the extra output allocation needed by that node.

node_memory(self, node_index: int) → dict[str, int | str | dict[str, int | str]]#: Returns the memory snapshot dict computed for the node at node_index. Raises IndexError when node_index is out of bounds.

node_release_after_shape_tagged(self, node_index: int) → list[str]#: Returns the list of shape-tagged releasable value names for the node at node_index. Raises IndexError when node_index is out of bounds.

node_reuse(self, node_index: int) → list[onnx_light.onnx_py._onnxpyoptim.shape_inference.InPlaceReuse]#: Returns the list of InPlaceReuse opportunities discovered for the node at node_index. Raises IndexError when node_index is out of bounds.

node_tag(self, node_index: int) → str#: Returns the last tag computed for the node at node_index. Raises IndexError when node_index is out of bounds.

property node_tags#: Returns the last per-node tag list computed by compute_value_and_node_tags().

property release_after_shape_tagged#: The per-node shape-tagged releasable values. When compute_inplace_reuse_graph was called with an explicit non-empty value_tags argument (or after compute_value_and_node_tags populated this context), this is a list with one entry per node (same order as graph.node), where each entry is a list of value names that carry the "shape" tag. Otherwise this list is itself empty.

property reuse#: The per-node reuse opportunities as a list (one entry per node, same order as graph.node); each entry is a list of InPlaceReuse.

property value_tags#: Returns the last value-tag map computed by compute_value_and_node_tags().

write_to_metadata(self, graph: onnx_light.onnx_py._onnxpyprotoop.GraphProto) → None#: Records the computed opportunities into each node’s metadata_props under the keys onnx_light.inplace_reuse, onnx_light.release_after, and (when value tags were provided to compute_inplace_reuse_graph) onnx_light.release_after_shape_tag. The GraphProto is mutated in place and must be the same graph passed to compute_inplace_reuse_graph.

class onnx_light.onnx_optim.shape_inference.ComputeEvent#

One entry in ComputeContext.events(). inplace events carry output_index / input_index / kind. release and release_shape_tag events carry name.

property action#: ComputeEventAction value describing the decision kind.

as_dict(self) → dict#: Returns this event as a plain dict.

property input_index#: Input index for inplace events, -1 otherwise.

property kind#: Reuse kind for inplace events.

property name#: Value name for release / release_shape_tag events.

property node_index#: Node index where the decision was made.

property output_index#: Output index for inplace events, -1 otherwise.

class onnx_light.onnx_optim.shape_inference.ComputeEventAction(*values)#

Classifies decisions logged by ComputeContext when events_enabled is True: kInPlace for in-place matches, kRelease for releasable last-use values, and kReleaseShapeTag for released values tagged "shape".

kInPlace = 0#

kRelease = 1#

kReleaseShapeTag = 2#

class onnx_light.onnx_optim.shape_inference.InPlaceReuse(*args, **kwargs)#

Represents one in-place reuse opportunity for a node: the output at output_index reuses the buffer of the input at input_index (both indices into the node’s output/input lists). kind records whether the input buffer has the same size as the output (kEqual) or is strictly larger (kGreater).

property input_index#: (self) -> int

property kind#: (self) -> onnx_light.onnx_py._onnxpyoptim.shape_inference.InPlaceReuseKind

property output_index#: (self) -> int

class onnx_light.onnx_optim.shape_inference.InPlaceReuseKind(*values)#

Classifies how the reused input buffer compares in size with the output: kEqual when the input and output share the same element type and shape (same byte size, the preferred reuse); kGreater when the input buffer is strictly larger in bytes than the output.

kEqual = 0#

kGreater = 1#

class onnx_light.onnx_optim.shape_inference.OptimDim(*args, **kwargs)#

A single shape dimension that is either a concrete integer or a symbolic string expression.

as_expr(self) → str#: Returns the symbolic expression. Raises if the dimension is an integer.

as_int(self) → int#: Returns the integer value. Raises if the dimension is symbolic.

is_expr(self) → bool#: Returns True when the dimension holds a symbolic string expression.

is_int(self) → bool#: Returns True when the dimension holds a concrete integer value.

value(self) → object#: Returns the underlying value as either int or str.

class onnx_light.onnx_optim.shape_inference.OptimShape(*args, **kwargs)#

Ordered, bounded-rank collection of OptimDim entries describing a tensor shape.

dims(self) → list#: Returns the dimensions as a list of int or str.

empty(self) → bool#: True when the shape is rank-0.

is_fully_known(self) → bool#: True when every dimension is a concrete integer.

rank(self) → int#: Number of dimensions.

class onnx_light.onnx_optim.shape_inference.OptimTensor(*args, **kwargs)#

Lightweight (non-owning) tensor descriptor with an element type and an OptimShape, optionally annotated with a value-as-shape and value bounds. The Python binding never references a buffer; only the descriptor metadata is carried.

clear_max(self) → None#

clear_min(self) → None#

clear_value_as_shape(self) → None#

property dtype#: Element type as a TensorProto.DataType integer.

has_max(self) → bool#

has_min(self) → bool#

has_value_as_shape(self) → bool#: True when the tensor’s value is interpreted as a shape.

is_null(self) → bool#: True when no data buffer is attached.

max(self) → float#

min(self) → float#

set_max(self, value: float) → None#

set_min(self, value: float) → None#

set_value_as_shape(self, shape: object) → None#: Tags the tensor as carrying a shape value (e.g. the shape input of Reshape) and stores that shape.

property shape#: Shape (a copy of the underlying OptimShape).

value_as_shape(self) → onnx_light.onnx_py._onnxpyoptim.shape_inference.OptimShape#: Returns the value-as-shape annotation. Raises if not set.

class onnx_light.onnx_optim.shape_inference.ShapeEvent#

One entry of the ShapesContext.events() log. add / replace events describe a tensor descriptor mutation performed through ShapesContext.set and carry the descriptor name, data_type (a TensorProto.DataType integer) and shape (a list of per-dimension strings, preserving symbolic dims). compute_node events summarise the shape-inference dispatch of a single node and carry op_domain, op_type and inputs instead. constraint / constraint_max events record a newly inserted symbolic-dimension constraint and carry its two operands in inputs.

property action#

kAdd, kReplace, kComputeNode, kConstraint or kConstraintMax.

Type:: ShapeEventAction member describing the event kind

as_dict(self) → dict#: Returns the event fields as a plain Python dict (trivially renderable as a table, serialisable, etc.).

property data_type#: TensorProto.DataType integer of the descriptor, or UNDEFINED (0) for compute_node / constraint / constraint_max events.

property inputs#

the two constraint operands. Empty otherwise.

Type:: For compute_node events
Type:: ordered list of input names consumed by the node (matching NodeProto.input). For constraint / constraint_max events

property name#: Value name targeted by the mutation. Empty for compute_node / constraint / constraint_max events.

property node_index#

-1 for graph inputs, -2 for initializers, and the position (>= 0) of the producing / dispatched node otherwise (-1 when no producing node is known).

Type:: Index of the node this event is associated with

property op_domain#

normalised ONNX op domain of the dispatched node (default domain reported as "ai.onnx"). Empty otherwise.

Type:: For compute_node events

property op_type#

ONNX op_type of the dispatched node. Empty otherwise.

Type:: For compute_node events

property shape#: Descriptor shape as a list of per-dimension strings (decimal integers for concrete dims, symbolic expressions otherwise). Empty for compute_node / constraint / constraint_max events.

property subgraph_attr_name#: Attribute name of the subgraph within the owning control-flow node ("body", "then_branch", "else_branch", etc.). Empty for top-level-graph events.

property subgraph_node_index#: Index of the control-flow node in the parent graph whose attribute subgraph produced this event. -1 for top-level-graph events.

class onnx_light.onnx_optim.shape_inference.ShapeEventAction(*values)#

Action kind recorded in a ShapeEvent. kAdd / kReplace mark tensor-descriptor mutations; kComputeNode marks the dispatch of a single shape-inference node; kConstraint / kConstraintMax mark newly inserted symbolic dimension constraints.

kAdd = 0#

kComputeNode = 2#

kConstraint = 3#

kConstraintMax = 4#

kReplace = 1#

class onnx_light.onnx_optim.shape_inference.ShapesContext(*args, **kwargs)#

In/out container shared by the per-operator ComputeShape* shape-inference functions. Holds a name -> OptimTensor map, a name -> OptimSequence map and a domain -> opset_version map mirroring opset_import.

add_constraint(self, a: str, b: str) → bool#: Records an equality constraint between two symbolic dimension names. The pair is canonicalised so (a, b) and (b, a) are stored only once, and a == a is dropped. Returns True when a new constraint was inserted, False otherwise.

add_less_equal_constraint(self, lhs: str, rhs: str) → bool#: Records that the symbolic dimension lhs is less than or equal to the dimension expression rhs. lhs == rhs and empty operands are dropped. Returns True when a new constraint was inserted, False otherwise.

apply_inferred_shapes_to_graph(self, graph: onnx_light.onnx_py._onnxpyprotoop.GraphProto) → None#: Writes the shape and element-type descriptors stored in self back into graph.output and graph.value_info.

apply_inferred_shapes_to_model(self, model: onnx_light.onnx_py._onnxpyprotoop.ModelProto) → None#: Writes the shape and element-type descriptors stored in self back into model.graph.

check_inputs_available(self, node: onnx_light.onnx_py._onnxpyprotoop.NodeProto) → None#: Raises ValueError if any non-empty input name declared by node is missing from self.

check_outputs_not_available(self, node: onnx_light.onnx_py._onnxpyprotoop.NodeProto) → None#: Raises ValueError if any non-empty output name declared by node already has an entry in self.

clear(self) → None#: Removes every entry (tensors, sequences and opset versions).

clear_custom_shape_inference_functions(self) → None#: Removes every registered custom callback.

clear_events(self) → None#: Empties the event log without otherwise touching the context.

compute_shape_graph(self, graph: onnx_light.onnx_py._onnxpyprotoop.GraphProto) → None#: Seeds self from the initializers and inputs of graph and then runs compute_shape_node on every node in topological order.

compute_shape_model(self, model: onnx_light.onnx_py._onnxpyprotoop.ModelProto, prefill_with_value_info_output: bool = False) → None#: Records every (domain, version) pair from model.opset_import in self and delegates to compute_shape_graph. When prefill_with_value_info_output is true, tensor descriptors from model.graph.value_info and model.graph.output are added as anchors and preferred when there is a non-conflicting choice at the end.

compute_shape_node(self, node: onnx_light.onnx_py._onnxpyprotoop.NodeProto) → None#: Dispatches a single NodeProto to the matching per-operator ComputeShape* function and stores the resulting output tensor descriptors in self. The node’s input descriptors must already be present in self.

constraints(self) → list#: Returns the list of recorded equality constraints as (lhs, rhs) tuples with lhs < rhs.

constraints_size(self) → int#: Number of recorded equality constraints.

custom_shape_inference_keys(self) → list#: Returns registered custom callback keys as "<domain>:<op_type>".

empty(self) → bool#: True when no entries are stored.

events(self) → list#: Returns the append-only shape-inference event log as a list of ShapeEvent instances. Empty unless events_enabled was set before running shape inference.

property events_enabled#: When True, set records add / replace events, compute_shape_node records a compute_node event per dispatched node, and add_constraint / add_less_equal_constraint record constraint / constraint_max events. Default is False for maximum throughput; enable only when tracing shape inference.

get(self, name: str) → onnx_light.onnx_py._onnxpyoptim.shape_inference.OptimTensor#: Returns (a copy of) the tensor descriptor stored under name. Raises KeyError when no such entry exists.

has(self, name: str) → bool#: True when a tensor descriptor is stored under name.

has_constraint(self, a: str, b: str) → bool#: True when an equality constraint between a and b was recorded (canonical order is applied before lookup; a == a always returns True).

has_custom_shape_inference_function(self, domain: str, op_type: str) → bool#: True when a custom callback is registered for (domain, op_type).

has_less_equal_constraint(self, lhs: str, rhs: str) → bool#: True when a lhs <= rhs constraint was recorded (lhs == rhs always returns True).

has_opset_version(self, domain: str) → bool#: True when an opset version has been recorded for domain.

has_sequence(self, name: str) → bool#: True when a sequence descriptor is stored under name.

has_subgraph_context(self, node_index: int, attr_name: str) → bool#: True when a child context was retained for the subgraph attr_name of the control-flow node at node_index.

less_equal_constraints(self) → list#: Returns the list of recorded <= constraints as ordered (lhs, rhs) tuples meaning lhs <= rhs.

less_equal_constraints_size(self) → int#: Number of recorded <= constraints.

names(self) → list#: Returns the list of names of stored tensor descriptors.

opset_version(self, domain: str) → int#: Returns the recorded opset version of domain, or kUnknownOpsetVersion (-1) when none was recorded.

opsets(self) → dict[str, int]#: Returns a copy of the domain -> opset_version map.

remove_custom_shape_inference_function(self, domain: str, op_type: str) → bool#: Removes a custom callback for (domain, op_type). An empty domain is normalised to ai.onnx. Returns True when an entry was removed.

sequence_names(self) → list#: Returns the list of names of stored sequence descriptors.

sequences_size(self) → int#: Number of sequence descriptors currently stored.

set(self, name: str, tensor: onnx_light.onnx_py._onnxpyoptim.shape_inference.OptimTensor) → None#: Inserts or replaces the tensor descriptor stored under name.

set_custom_shape_inference_function(self, domain: str, op_type: str, fn: collections.abc.Callable) → None#: Registers a Python callback for (domain, op_type). The callback receives (ctx, node) and can populate output descriptors in ctx. An empty domain is normalised to ai.onnx.

set_opset_version(self, domain: str, opset_version: int) → None#: Records the opset version for domain. An empty domain is normalised to ai.onnx.

size(self) → int#: Number of tensor descriptors currently stored.

subgraph_context(self, node_index: int, attr_name: str) → onnx_light.onnx_py._onnxpyoptim.shape_inference.ShapesContext#: Returns the child ShapesContext retained for the subgraph attr_name of the control-flow node at node_index. Raises IndexError if absent.

subgraph_context_keys(self) → list#: Returns retained child-context keys as (node_index, attr_name) tuples.

subgraph_contexts_size(self) → int#: Number of retained child contexts.

onnx_light.onnx_optim.shape_inference.infer_shapes_model(model, prefill_with_value_info_output: bool = False) → None#

Runs shape inference on a ModelProto in place.

Seeds the shape-inference context from the model’s opset_import, graph initializers and graph inputs, runs per-operator shape inference on every node of the main graph, and writes the inferred element types and shapes back into model.graph.output and model.graph.value_info.

Parameters:

model – A ModelProto to mutate in place.
prefill_with_value_info_output – When True, prefill from model.graph.value_info and model.graph.output and prefer these anchors when there is a non-conflicting end-state alternative.

Raises:

ValueError – If shape inference rejects a node (for example because of an unsupported op type) or if model has no graph.