Implementation:Onnx Onnx Shape Inference Engine
| Knowledge Sources | |
|---|---|
| Domains | Shape Inference, Graph Traversal, Data Propagation |
| Last Updated | 2026-02-10 00:00 GMT |
Overview
Concrete tool for providing the internal implementation data structures and concrete classes that drive the ONNX shape inference engine, provided by the ONNX library.
Description
The implementation.h header defines the concrete implementation classes that power the ONNX shape inference system. While shape_inference.h defines the abstract interfaces, this header provides the working implementations used by the inference engine.
SymbolTableImpl implements the SymbolTable interface for managing symbolic dimension names. It traverses graph inputs, outputs, and value_info to collect existing symbolic dimension parameters (e.g., "batch_size", "seq_len"), then generates new unique symbols with a counter-based scheme that avoids collisions. The symbol table handles all ONNX type categories: tensors, sparse tensors, sequences, optionals, and maps.
GraphInferenceContext is a struct that bundles all the state needed for graph-level inference: outer scope value types (for nested graphs), opset imports, a symbol table reference, model-local function definitions, a schema registry, generated shape data, and the IR version. This context is threaded through the inference process to provide access to global state.
GraphInferencerImpl implements GraphInferencer to perform inference on subgraphs contained within control-flow operators (e.g., If, Loop). It takes a mutable GraphProto and a GraphInferenceContext, and produces inferred output types based on provided input types and data.
InferenceContextImpl is the main concrete implementation of InferenceContext. It is constructed from a NodeProto and several name-to-type/data maps, populating internal vectors of input types, input data (from TensorProto initializers, SparseTensorProto, or generated shape data), and output types. It provides methods for accessing attributes by name, creating GraphInferencerImpl instances on demand for graph-valued attributes, and generating display names for error reporting.
DataPropagationContextImpl implements DataPropagationContext for partial evaluation of tensor values. It supports converting initializer data (int32/int64 scalars and 1D tensors) into TensorShapeProto representations, storing and retrieving generated shape data, and inferring value arrays from known tensor ranks. This enables the shape inference system to track tensor values through operations like Shape, Slice, and Concat to determine shapes for downstream Reshape-like operators.
The file also declares the top-level InferShapes API functions that accept GraphProto, ModelProto, or file paths, along with InferShapeForFunctionNode for function-level inference and InferFunctionOutputTypes for inferring function output types from input types and attributes.
Usage
This header is used internally by the ONNX shape inference engine and by tools that need low-level access to inference contexts. The InferShapes functions are the primary entry points for performing shape inference on models, while the context classes are used when implementing custom inference logic or extending the inference system.
Code Reference
Source Location
- Repository: Onnx_Onnx
- File: onnx/shape_inference/implementation.h
- Lines: 1-545
Signature
// Type aliases
using ModelLocalFunctionsMap = std::unordered_map<std::string, const FunctionProto*>;
using DataValueMap = std::unordered_map<std::string, TensorShapeProto>;
// Symbol table implementation
class SymbolTableImpl : public SymbolTable {
public:
void addFromGraph(const GraphProto& g) override;
std::string createNew(const std::string& symbol_prefix) override;
};
// Graph-level inference context
struct GraphInferenceContext {
GraphInferenceContext(
const std::unordered_map<std::string, TypeProto*>& outer_scope_value_types_by_name_in,
std::unordered_map<std::string, int> opset_imports_in,
SymbolTable* symbol_table_in = nullptr,
const ModelLocalFunctionsMap& model_local_functions_in = {},
const ISchemaRegistry* schema_registry_in = OpSchemaRegistry::Instance(),
DataValueMap* generated_shape_data_by_name_in = nullptr,
const int ir_version_in = IR_VERSION);
};
// Node-level inference context
struct InferenceContextImpl : public InferenceContext {
InferenceContextImpl(
NodeProto& n,
const std::unordered_map<std::string, TypeProto*>& valueTypesByName,
const std::unordered_map<std::string, const TensorProto*>& inputDataByName,
const std::unordered_map<std::string, const SparseTensorProto*>& inputSparseDataByName,
const ShapeInferenceOptions& options,
DataValueMap* generatedShapeData = nullptr,
GraphInferenceContext* graphInferenceContext = nullptr);
};
// Top-level InferShapes API
void InferShapes(ModelProto& m,
const ISchemaRegistry* schema_registry = OpSchemaRegistry::Instance(),
const ShapeInferenceOptions& options = ShapeInferenceOptions(),
DataValueMap* generated_shape_data_by_name = nullptr);
void InferShapes(GraphProto* g,
const std::unordered_map<std::string, int>& opset_imports,
const ISchemaRegistry* schema_registry = OpSchemaRegistry::Instance(),
const ShapeInferenceOptions& options = ShapeInferenceOptions(),
const ModelLocalFunctionsMap& in_model_functions = {});
void InferShapes(const std::string& model_path,
const std::string& save_path = "",
const ISchemaRegistry* schema_registry = OpSchemaRegistry::Instance(),
const ShapeInferenceOptions& options = ShapeInferenceOptions(),
DataValueMap* generated_shape_data_by_name = nullptr);
std::vector<TypeProto> InferFunctionOutputTypes(
const FunctionProto& function_proto,
const std::vector<TypeProto>& input_types,
const std::vector<AttributeProto>& attributes);
Import
#include "onnx/shape_inference/implementation.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| m | ModelProto& | Yes (model overload) | The ONNX model to perform shape inference on; modified in-place |
| g | GraphProto* | Yes (graph overload) | The graph to perform shape inference on; modified in-place |
| model_path | const std::string& | Yes (path overload) | Path to the ONNX model file |
| save_path | const std::string& | No | Path to save the inferred model; empty means no save |
| schema_registry | const ISchemaRegistry* | No | Schema registry for operator lookups; defaults to OpSchemaRegistry::Instance() |
| options | const ShapeInferenceOptions& | No | Inference options controlling type checking, error mode, and data propagation |
| opset_imports | std::unordered_map<std::string, int> | Yes (graph overload) | Map of domain to opset version |
| generated_shape_data_by_name | DataValueMap* | No | Map to store/retrieve statically computed tensor values |
| function_proto | const FunctionProto& | Yes (function overload) | The function to infer output types for |
| input_types | std::vector<TypeProto> | Yes (InferFunctionOutputTypes) | Types of the function inputs |
| attributes | std::vector<AttributeProto> | Yes (InferFunctionOutputTypes) | Attribute values supplied to the function |
Outputs
| Name | Type | Description |
|---|---|---|
| (modified model/graph) | ModelProto& / GraphProto* | The model or graph is modified in-place with inferred type and shape information on all value_info entries |
| output_types | std::vector<TypeProto> | Inferred output types of a function (for InferFunctionOutputTypes) |
| (generated data) | DataValueMap* | Populated with statically known tensor values discovered during inference |
Usage Examples
#include "onnx/shape_inference/implementation.h"
using namespace ONNX_NAMESPACE;
// Infer shapes on a model
ModelProto model;
// ... load model ...
shape_inference::InferShapes(model);
// Infer shapes with strict mode and data propagation
ShapeInferenceOptions options(/*check_type=*/true, /*strict_mode=*/1, /*data_prop=*/true);
shape_inference::InferShapes(model, OpSchemaRegistry::Instance(), options);
// Infer shapes with data value tracking
shape_inference::DataValueMap generated_data;
shape_inference::InferShapes(model, OpSchemaRegistry::Instance(),
ShapeInferenceOptions(), &generated_data);
// Infer shapes on a file
shape_inference::InferShapes("input.onnx", "output.onnx");
// Infer function output types
FunctionProto func;
std::vector<TypeProto> input_types = { /* ... */ };
std::vector<AttributeProto> attrs = { /* ... */ };
auto output_types = shape_inference::InferFunctionOutputTypes(func, input_types, attrs);
// Using SymbolTableImpl
shape_inference::SymbolTableImpl symbol_table;
symbol_table.addFromGraph(model.graph());
std::string new_sym = symbol_table.createNew("batch_");