Implementation:Onnx Onnx Shape Inference Utils
| Knowledge Sources | |
|---|---|
| Domains | Shape Inference, Type Propagation, Graph Analysis |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
Concrete tool for propagating and merging type and shape information across ONNX graph nodes during static analysis, provided by the ONNX library.
Description
The shape_inference.cc file implements the core shape inference utility functions used throughout the ONNX framework to statically determine tensor types and shapes at every point in the computation graph. These functions are responsible for propagating element types from inputs to outputs, merging shape information between TypeProto instances, and performing union operations on shapes.
The file handles all ONNX type categories including dense tensors, sparse tensors, sequences, optionals, and maps. Each type category has dedicated propagation functions (e.g., propagateElemTypeFromTensorInputToOutput, propagateElemTypeFromSequenceInputToOutput) that correctly handle the type-specific protobuf structures. These are unified via the dispatch function propagateElemTypeFromInputToOutput which routes to the appropriate handler based on the input's value case.
Shape merging is performed by the family of mergeInShapeInfo overloads, which combine shape information from a source into a target following strict rules: values are preferred over symbolic parameters, dimension values must match when both sides specify them, and dimension counts must be equal. The UnionShapeInfo and UnionTypeInfo functions compute the least-specific shape that is consistent with both inputs, which is useful for control-flow operators where different branches may produce different shapes.
The file also provides getShapeInput which extracts shape values from initializers, symbolic inputs, or infers rank from shape-of-shape inputs. The getAttributeElementTypeAndLength utility inspects operator attributes to determine element types and lengths for validation purposes.
Usage
These utilities are used by operator-specific shape inference functions registered in ONNX operator schemas. They are called during model validation, optimization, and prior to execution to determine output tensor shapes from known input shapes. Runtime implementations and model compilers rely on these utilities to perform memory planning and graph optimization.
Code Reference
Source Location
- Repository: Onnx_Onnx
- File: onnx/defs/shape_inference.cc
- Lines: 1-559
Signature
void propagateElemTypeFromTensorInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
void propagateElemTypeFromInputToOutput(InferenceContext& ctx, size_t inputIndex, size_t outputIndex);
void mergeInShapeInfo(const TensorShapeProto& source_shape, TypeProto_Tensor& target_type);
void mergeInShapeInfo(const TensorShapeProto& source_shape, TypeProto_SparseTensor& target_type);
void mergeInShapeInfo(const TypeProto_Tensor& source, TypeProto_Tensor& target);
void mergeInShapeInfo(const TypeProto_SparseTensor& source, TypeProto_SparseTensor& target);
void UnionShapeInfo(const TensorShapeProto& source_shape, TypeProto_Tensor& target_type);
void UnionShapeInfo(const TensorShapeProto& source_shape, TypeProto_SparseTensor& target_type);
void UnionTypeInfo(const TypeProto& source_type, TypeProto& target_type);
void propagateElemTypeWithValidation(const TypeProto* input_type, TypeProto* output_type);
TensorShapeProto getShapeInput(const InferenceContext& ctx, size_t input_index, bool& found);
TensorShapeProto getShapeInput(const InferenceContext& ctx, size_t input_index, bool fail_if_negative_value, bool& found);
std::pair<int, int> getAttributeProtoElemTypeAndLength(const AttributeProto* attr_proto);
std::pair<int, int> getAttributeElementTypeAndLength(
const InferenceContext& ctx,
const std::initializer_list<std::string>& attribute_names);
Import
#include "onnx/defs/shape_inference.h"
#include "onnx/defs/data_type_utils.h"
#include "onnx/defs/tensor_proto_util.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| ctx | InferenceContext& | Yes | The inference context providing access to node inputs, outputs, and attributes |
| inputIndex | size_t | Yes | Index of the input whose type/shape is being propagated |
| outputIndex | size_t | Yes | Index of the output to receive the propagated type/shape |
| source_shape | const TensorShapeProto& | Varies | Source shape information to merge or union |
| source_type | const TypeProto& | Varies | Source type information to merge or union |
| input_type | const TypeProto* | Varies | Input type pointer for validated propagation |
Outputs
| Name | Type | Description |
|---|---|---|
| (modified output type) | TypeProto* | Output type is modified in-place with propagated element type and shape |
| (modified target type) | TypeProto_Tensor& / TypeProto_SparseTensor& | Target type is modified in-place with merged or unioned shape information |
| shape_input | TensorShapeProto | Shape values extracted from initializer, symbolic input, or rank inference |
| (elem_type, length) | std::pair<int, int> | Element type and length extracted from attribute protos |
Usage Examples
// Propagating element type from first input to first output
void MyOpInference(InferenceContext& ctx) {
propagateElemTypeFromTensorInputToOutput(ctx, 0, 0);
}
// Using propagateElemTypeFromInputToOutput (handles all type cases)
void MyGenericOpInference(InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
}
// Merging shape information from a known shape into an output type
void MergeExample(const TensorShapeProto& known_shape, TypeProto_Tensor& output_type) {
mergeInShapeInfo(known_shape, output_type);
}
// Extracting a shape input (e.g., the "shape" input of a Reshape operator)
void ReshapeInference(InferenceContext& ctx) {
bool found = false;
TensorShapeProto shape = getShapeInput(ctx, 1, found);
if (found) {
// Use shape to set output dimensions
}
}
// Propagating type with validation (checks for consistency)
void ValidatedPropagation(const TypeProto* input, TypeProto* output) {
propagateElemTypeWithValidation(input, output);
}