shapes#
The shapes sub-namespace of onnx_core (core::shapes) hosts the
generic shape-inference engine: ShapesContext, the node/graph
traversal (ComputeShapeNode, ComputeShapeGraph,
InferShapesModel()), broadcasting and node-checking helpers, and
the dispatch table that maps an operator (domain, op_type) pair to the
function that computes its output shapes.
onnx_core never depends on onnx_shapes, so the dispatch table starts
out empty: it is a mutable registry
(RegisterComputeShapeFn()) that onnx_shapes populates with its
per-operator ComputeShape* functions (see
dispatch_table.h) via
onnx_light::onnx_shapes::RegisterShapeFunctions(). Any consumer of
the shape-inference engine (Python bindings, tests, examples, …) must call
that function once before using InferShapesModel() or
ShapesContext.
Peak-memory estimation#
A second, parallel registry estimates each operator’s peak computation
memory rather than its output shapes. Mirroring the shape dispatch table, it
maps an (domain, op_type, device) identifier to a
core::shapes::ComputePeakMemoryFn — a function that takes the
Device the operator runs on followed by the
SymShape of each input and returns the estimated scratch memory
in bytes. Functions are registered with
RegisterComputePeakMemoryFn() and looked up through
ComputePeakMemory(); operators without a registered function report
0 by default. onnx_shapes populates the built-in estimators (for
example Attention) via
onnx_light::onnx_shapes::RegisterPeakMemoryFunctions(), and the
Python bindings expose compute_peak_memory() together with the
Device enum.