Implementation:Onnx Onnx Shape Inference Interfaces
| Knowledge Sources | |
|---|---|
| Domains | Shape Inference, Type System, Inference Framework |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
Concrete tool for defining the interfaces, types, and inline helper functions that constitute the ONNX shape and type inference framework, provided by the ONNX library.
Description
The shape_inference.h header file defines the core abstractions and utilities for the ONNX shape inference system. It is the central header that every operator schema includes when defining its shape inference behavior.
The file declares several key abstract interfaces. InferenceContext provides the contract for accessing node attributes, input types, input data (initializers), output types, symbolic inputs, and graph attribute inferencers during shape inference. DataPropagationContext is a related interface used for partial evaluation of tensor values to improve shape inference precision. GraphInferencer enables inference on subgraphs contained in control-flow operators like If and Loop. SymbolTable manages symbolic dimension names to maintain consistency across the graph.
The ShapeInferenceOptions struct configures the inference engine with options for type checking, error mode control (strict vs. lenient), and data propagation enablement. The InferenceError exception class provides structured error reporting with context tracking for debugging.
The file includes a comprehensive set of inline helper functions for common inference operations: propagateElemTypeFromDtypeToOutput sets output types from data type constants; propagateShapeFromInputToOutput copies shape information between inputs and outputs; updateOutputElemType and updateOutputShape modify output type information; multidirectionalBroadcastShapeInference and bidirectionalBroadcastShapeInference compute broadcast output shapes; mergeInDimensionInfo combines dimension information with conflict detection; and unifyDim / unifyInputDim perform dimension unification with constraint checking.
The header also defines type aliases InferenceFunction and DataPropagationFunction as std::function types, along with dummyInferenceFunction and dummyDataPropagationFunction no-op placeholders for operators without inference implementations.
Usage
This header is included by every ONNX operator schema definition that implements shape inference. It is also used by inference engine implementations, model validation tools, and custom operator authors who need to define shape inference behavior for their operators. The inline helper functions provide the building blocks that operator-specific inference functions compose to propagate type and shape information.
Code Reference
Source Location
- Repository: Onnx_Onnx
- File: onnx/defs/shape_inference.h
- Lines: 1-958
Signature
// Core abstract interfaces
struct InferenceContext {
virtual const AttributeProto* getAttribute(const std::string& name) const = 0;
virtual size_t getNumInputs() const = 0;
virtual const TypeProto* getInputType(size_t index) const = 0;
virtual const TensorProto* getInputData(size_t index) const = 0;
virtual size_t getNumOutputs() const = 0;
virtual TypeProto* getOutputType(size_t index) = 0;
virtual GraphInferencer* getGraphAttributeInferencer(const std::string& attribute_name) = 0;
virtual const SparseTensorProto* getInputSparseData(size_t index) const = 0;
virtual const TensorShapeProto* getSymbolicInput(size_t index) const = 0;
};
struct DataPropagationContext {
virtual const AttributeProto* getAttribute(const std::string& name) const = 0;
virtual size_t getNumInputs() const = 0;
virtual const TypeProto* getInputType(size_t index) const = 0;
virtual const TensorShapeProto* getInputData(size_t index) = 0;
virtual void addOutputData(size_t index, TensorShapeProto&& tp) = 0;
};
struct ShapeInferenceOptions {
bool check_type;
int error_mode;
bool enable_data_propagation;
};
// Function type aliases
using InferenceFunction = std::function<void(InferenceContext&)>;
using DataPropagationFunction = std::function<void(DataPropagationContext&)>;
// Key inline helper function declarations
void propagateShapeAndTypeFromFirstInput(InferenceContext& ctx);
void propagateElemTypeFromAttributeToOutput(InferenceContext& ctx, const std::string& attributeName, size_t outputIndex);
void multidirectionalBroadcastShapeInference(const std::vector<const TensorShapeProto*>& shapes, TensorShapeProto& resultShape);
void bidirectionalBroadcastShapeInference(const TensorShapeProto& shapeL, const TensorShapeProto& shapeR, TensorShapeProto& resultShape);
// Exported shape inference functions for common operator patterns
void RNNShapeInference(InferenceContext& ctx);
void convPoolShapeInference(InferenceContext& ctx, bool use_dilation, bool require_kernel_shape, int input1Idx, int input2Idx);
void convTransposeShapeInference(InferenceContext& ctx);
void globalPoolTypeShapeInference(InferenceContext& ctx);
Import
#include "onnx/defs/shape_inference.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| ctx | InferenceContext& | Yes | The inference context providing access to node inputs, outputs, and attributes |
| shapes | std::vector<const TensorShapeProto*> | Varies | Input shapes for broadcast inference |
| attributeName | const std::string& | Varies | Name of attribute to read type/shape information from |
| dim1 / dim2 | const TensorShapeProto::Dimension& | Varies | Dimensions for arithmetic or unification operations |
| input_index / output_index | size_t | Varies | Indices of inputs or outputs to access |
Outputs
| Name | Type | Description |
|---|---|---|
| (modified output type) | TypeProto* | Output type modified in-place through the InferenceContext |
| resultShape | TensorShapeProto& | The computed broadcast shape, modified in-place |
| (dimension result) | TensorShapeProto::Dimension | Result of dimension arithmetic operations (operator*, operator/) |
| (exception) | InferenceError | Thrown on type or shape inference failures with context tracking |
Usage Examples
// Simple shape/type pass-through inference (e.g., for Relu, Sigmoid)
ONNX_OPERATOR_SET_SCHEMA(Relu, 14,
OpSchema()
.SetDoc("Relu activation")
.Input(0, "X", "Input tensor", "T")
.Output(0, "Y", "Output tensor", "T")
.TypeAndShapeInferenceFunction(propagateShapeAndTypeFromFirstInput));
// Custom inference with broadcast
void MyAddInference(InferenceContext& ctx) {
propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (hasNInputShapes(ctx, 2)) {
bidirectionalBroadcastShapeInference(
getInputShape(ctx, 0),
getInputShape(ctx, 1),
*getOutputShape(ctx, 0));
}
}
// Reading an attribute to determine output type
void CastInference(InferenceContext& ctx) {
propagateElemTypeFromAttributeToOutput(ctx, "to", 0);
if (hasInputShape(ctx, 0)) {
propagateShapeFromInputToOutput(ctx, 0, 0);
}
}
// Checking input rank
void MatMulInference(InferenceContext& ctx) {
checkInputRank(ctx, 0, 2);
checkInputRank(ctx, 1, 2);
}
// Dimension unification
void ConcatInference(InferenceContext& ctx) {
Dim batch_dim;
unifyInputDim(ctx, 0, 0, batch_dim);
unifyInputDim(ctx, 1, 0, batch_dim);
}