shapes_context.h#

Name → :cpp:class:OptimTensor map shared by all onnx_optim shape-inference functions.

ShapesContext is the in/out parameter consumed and produced by the per-operator ComputeShape* functions (for example :cpp:func:ComputeShapeAbs). It holds the :cpp:class:OptimTensor descriptors of every named value (graph input, initializer or intermediate result) currently known to a shape-inference pass. ComputeShape* functions read the entries corresponding to a node’s inputs and insert new entries for the node’s outputs.

namespace onnx_light

Alias that makes onnx-light headers compatible with code that references ONNX_LIGHT_NAMESPACE (the macro used in the standard onnx package).

Set to ONNX_LIGHT_NAMESPACE so both names resolve to the same namespace.

Symbol-visibility attribute for the public onnx-light C++ API.

Defined as empty because onnx-light does not require explicit __declspec(dllexport) or __attribute__((visibility("default"))) annotations — visibility is controlled at the shared-library level. The macro is provided so that vendored ONNX headers that decorate their declarations with ONNX_API compile without modification.

Namespace alias so that ONNX C++ code (and consumers such as onnxruntime) that refers to the literal onnx namespace — rather than the ONNX_NAMESPACE macro — resolves to the onnx-light namespace. The standard onnx package lives in namespace onnx; onnx-light uses onnx_light (via ONNX_LIGHT_NAMESPACE), so this alias keeps onnx-light a true drop-in. It is only introduced when the onnx-light namespace differs from onnx.

namespace onnx_optim

namespace shapes

Typedefs

using ShapeEventLog = std::vector<ShapeEvent>#: Append-only log of shape-inference events recorded by :cpp:class:ShapesContext.

Enums

enum class ShapeEventAction : int32_t#

Kind of shape-inference event recorded in the :cpp:class:ShapesContext event log. Mirrors the runtime :cpp:enum:onnx_kernels::RuntimeEventAction design.

kAdd — a new tensor descriptor was inserted via :cpp:func:ShapesContext::Set on a previously absent name.
kReplace — an existing tensor descriptor was overwritten via :cpp:func:ShapesContext::Set.
kComputeNode — shape inference was dispatched for a single :cpp:class:NodeProto. The event records the node’s op_domain / op_type and the list of inputs it consumed. It does not mutate the tensor map by itself.
kConstraint — a new symbolic-dimension equality constraint (a == b) was recorded via :cpp:func:ShapesContext::AddConstraint. The two operands are stored in inputs.
kConstraintMax — a new symbolic-dimension upper-bound constraint (lhs <= rhs) was recorded via :cpp:func:ShapesContext::AddLessEqualConstraint. The two operands are stored in inputs.

Values:

enumerator kAdd#

enumerator kReplace#

enumerator kComputeNode#

enumerator kConstraint#

enumerator kConstraintMax#

Functions

const char *ShapeEventActionName(ShapeEventAction action) noexcept#: Returns a short lowercase label for action ("add", "replace", "compute_node"). Useful for human-readable rendering of the event log.

Variables

constexpr int kUnknownOpsetVersion = -1#: Sentinel value returned by :cpp:func:ShapesContext::OpsetVersion when no opset version has been recorded for the requested domain.

constexpr const char *kOnnxDomain = "ai.onnx"#: Canonical domain string used for the standard ONNX operator set (ai.onnx). An empty domain on a NodeProto is treated as equivalent to this value.

struct ShapeEvent#

#include <shapes_context.h>

Single entry of the :cpp:class:ShapesContext event log.

Each insertion or replacement of a tensor descriptor performed through :cpp:func:ShapesContext::Set produces one ShapeEvent capturing the action, the name of the value and a snapshot of the descriptor’s element type and shape. Because shape inference works on descriptors rather than data, no element values are captured; the shape is recorded as a list of per-dimension strings so symbolic dimensions ("N", "2*N") are preserved alongside concrete integer dimensions.

:cpp:enumerator:ShapeEventAction::kComputeNode events instead summarise the dispatch of a single NodeProto: they carry the node’s op_domain / op_type and the list of inputs consumed; data_type is set to DataType::UNDEFINED and shape is left empty.

:cpp:enumerator:ShapeEventAction::kConstraint / :cpp:enumerator:ShapeEventAction::kConstraintMax events record a newly inserted symbolic-dimension constraint; the two operands are stored in inputs ({a, b} for an equality a == b, {lhs, rhs} for an upper bound lhs <= rhs).

Public Members

ShapeEventAction action = ShapeEventAction::kAdd #: Kind of event recorded by this entry.

std::string name#: Name of the value targeted by the event.

int32_t data_type = 0#: Element data type of the descriptor at the moment of the event, encoded as a TensorProto::DataType integer value. Set to DataType::UNDEFINED for kComputeNode / kConstraint / kConstraintMax events.

std::vector<std::string> shape#: Descriptor shape at the moment of the event, with each dimension rendered as a string (a decimal integer for concrete dims, the symbolic expression otherwise). Empty for kComputeNode / kConstraint / kConstraintMax events.

std::string op_domain#: For kComputeNode events: ONNX op domain of the dispatched node, normalised so the default domain is reported as "ai.onnx". Empty for other events.

std::string op_type#: For kComputeNode events: ONNX op_type of the dispatched node. Empty for other events.

std::vector<std::string> inputs#: For kComputeNode events: ordered list of input names consumed by the node, matching NodeProto::input. For kConstraint / kConstraintMax events: the two constraint operands ({a, b} or {lhs, rhs}). Empty for kAdd / kReplace events.

int64_t node_index = -1#: Index of the node this event is associated with. For graph inputs it is -1 and for initializers it is -2. For intermediate / output descriptors and for kComputeNode / kConstraint / kConstraintMax events it is the position (>= 0) of the producing / dispatched node in its graph node list. -1 when no producing node is known.

int64_t subgraph_node_index = -1#: Index of the control-flow node in the parent graph whose attribute subgraph produced this event. -1 for events from the top-level graph. Combined with :cpp:var:subgraph_attr_name this uniquely identifies which operator and which attribute subgraph an event originated from.

std::string subgraph_attr_name#: Attribute name of the subgraph within the control-flow node identified by :cpp:var:subgraph_node_index: "body" for :onnx:Loop / :onnx:Scan, "then_branch" or "else_branch" for :onnx:If. Empty for top-level-graph events.

class ShapesContext#

#include <shapes_context.h>

Lightweight container shared by the per-operator ComputeShape* shape-inference functions. ShapesContext carries two pieces of information:

a name → OptimTensor map describing every named value (graph input, initializer or intermediate result) currently known to the shape-inference pass;
a name → OptimSequence map describing every named sequence-typed value (the output of SequenceConstruct, SequenceEmpty, SplitToSequence, …);
a domain → opset_version map mirroring the opset_import entries of the surrounding ModelProto, so that ComputeShape* functions can pick the correct schema revision when shape inference depends on the operator’s opset version.

The context is a thin wrapper and does not own any tensor data: the :cpp:class:OptimTensor values stored here are themselves non-owning views.

Public Types

using CustomComputeShapeFn = std::function<void(ShapesContext&, const NodeProto&)>#

using CustomShapeInferenceMap = std::unordered_map<std::string, CustomComputeShapeFn>#

using SubgraphContextKey = std::pair<int64_t, std::string>#: Key identifying a child :cpp:class:ShapesContext retained while inferring a control-flow node’s attribute subgraph: the index of the control-flow node in this context’s graph paired with the name of the attribute carrying the subgraph ("body" for :onnx:Loop / :onnx:Scan, "then_branch" / "else_branch" for :onnx:If).

using Constraint = std::pair<std::string, std::string>#: Type used to store a single symbolic equality constraint.

using LessEqualConstraint = std::pair<std::string, std::string>#: Type used to store a single symbolic <= (less-or-equal) upper-bound constraint. The first element is the bounded symbol, the second is an arbitrary dimension expression that upper-bounds it.

Public Functions

ShapesContext() = default#

inline void Set(const std::string &name, OptimTensor &&tensor)#: Inserts or replaces the descriptor for name. tensor is consumed; callers must pass an rvalue (use std::move). When event logging is enabled (see :cpp:func:set_events_enabled) a :cpp:class:ShapeEvent describing the new state is appended to the event log — with action :cpp:enumerator:ShapeEventAction::kAdd when name was absent and :cpp:enumerator:ShapeEventAction::kReplace otherwise.

inline void Set(const char *name, OptimTensor &&tensor)#: Overload: name given as a null-terminated C string.

inline void Set(const utils::String &name, OptimTensor &&tensor)#: Overload: name given as a :cpp:class:utils::String.

inline bool Has(const std::string &name) const#: Returns true when an entry exists for name.

inline bool Has(const utils::String &name) const#: Overload: name given as a :cpp:class:utils::String.

inline const OptimTensor &Get(const std::string &name) const#: Returns the descriptor for name. Throws std::out_of_range if no such entry exists.

inline const OptimTensor &Get(const utils::String &name) const#: Overload: name given as a :cpp:class:utils::String.

inline std::size_t Size() const noexcept#: Number of named entries currently stored.

inline bool Empty() const noexcept#: true when no entries are stored.

inline void Clear() noexcept#: Removes every entry (both tensor descriptors and opset versions).

inline const std::string &TopKKDimName(const std::string &k_input_name)#: Returns (and records) the symbolic dimension name to use for the TopK output axis driven by the K input named k_input_name. The first unique K input seen returns "TopK_k"; each subsequent distinct K input gets "TopK_k_2", "TopK_k_3", and so on. Calling this method twice with the same k_input_name always returns the same string.

inline const std::unordered_map<std::string, OptimTensor> &Tensors() const noexcept#: Read-only access to the underlying map (useful for iteration).

inline void SetSequence(const std::string &name, OptimSequence &&sequence)#: Inserts or replaces the descriptor for a sequence-typed value named name. sequence is consumed; callers must pass an rvalue (use std::move).

inline void SetSequence(const char *name, OptimSequence &&sequence)#: Overload: name given as a null-terminated C string.

inline void SetSequence(const utils::String &name, OptimSequence &&sequence)#: Overload: name given as a :cpp:class:utils::String.

inline bool HasSequence(const std::string &name) const#: Returns true when a sequence-typed entry exists for name.

inline bool HasSequence(const utils::String &name) const#: Overload: name given as a :cpp:class:utils::String.

inline const OptimSequence &GetSequence(const std::string &name) const#: Returns the sequence descriptor for name. Throws std::out_of_range if no such entry exists.

inline const OptimSequence &GetSequence(const utils::String &name) const#: Overload: name given as a :cpp:class:utils::String.

inline std::size_t SequencesSize() const noexcept#: Number of sequence-typed entries currently stored.

inline const std::unordered_map<std::string, OptimSequence> &Sequences() const noexcept#: Read-only access to the underlying sequence map (useful for iteration).

void RegisterSubgraphContext(int64_t node_index, const std::string &attr_name, ShapesContext context)#: Retains the child context context produced while inferring the subgraph attr_name of the control-flow node at node_index so that the subgraph’s internal descriptors stay inspectable once the parent inference has completed. context is consumed (moved into the store). Any context previously registered for the same key is replaced.

inline bool HasSubgraphContext(int64_t node_index, const std::string &attr_name) const#: Returns true when a child context was registered for the subgraph attr_name of the control-flow node at node_index.

inline const ShapesContext &GetSubgraphContext(int64_t node_index, const std::string &attr_name) const#: Returns the child context registered for the subgraph attr_name of the control-flow node at node_index. Throws std::out_of_range if no such context exists.

inline std::size_t SubgraphContextsSize() const noexcept#: Number of retained child contexts.

inline const std::map<SubgraphContextKey, std::shared_ptr<ShapesContext>> &SubgraphContexts() const noexcept#: Read-only access to the retained child-context map (useful for iteration).

inline void ClearSubgraphContexts() noexcept#: Empties the retained child-context map without modifying other context state (tensor / sequence descriptors, opsets, constraints, events, …).

inline void SetOpsetVersion(const std::string &domain, int opset_version)#: Records the opset version of domain. An empty domain is normalised to :cpp:var:kOnnxDomain. Replaces any previous entry for the same domain.

inline bool HasOpsetVersion(const std::string &domain) const#: true when an opset version has been recorded for domain (after normalising the empty domain to :cpp:var:kOnnxDomain).

inline int OpsetVersion(const std::string &domain) const#: Returns the recorded opset version of domain, or :cpp:var:kUnknownOpsetVersion if none was recorded. An empty domain is normalised to :cpp:var:kOnnxDomain.

inline const std::unordered_map<std::string, int> &Opsets() const noexcept#: Read-only access to the underlying domain → opset_version map.

inline void SetLocalFunction(const FunctionProto *func)#: Registers a non-owning pointer to a model-local :cpp:class:FunctionProto. The pointer must remain valid for the lifetime of the shape-inference pass. Replaces any previous entry registered under the same "<domain>:<name>" key. func must not be nullptr.

inline bool HasLocalFunction(const std::string &key) const#: true when a model-local function is registered for key. key is expected to be the "<domain>:<name>" identifier of the function.

inline const FunctionProto *GetLocalFunction(const std::string &key) const#: Returns the registered FunctionProto pointer for key, or nullptr when none is registered. key is expected to be the "<domain>:<name>" identifier of the function.

inline const std::unordered_map<std::string, const FunctionProto*> &LocalFunctions() const noexcept#: Read-only access to the underlying "<domain>:<name>" → FunctionProto* map.

inline void SetCustomShapeInferenceFunction(const std::string &domain, const std::string &op_type, CustomComputeShapeFn fn)#: Registers or replaces a custom shape-inference callback for the (domain, op_type) pair. domain == "" is normalized to :cpp:var:kOnnxDomain.

inline const CustomComputeShapeFn *GetCustomShapeInferenceFunction(const std::string &domain, const std::string &op_type) const#: Returns a pointer to the custom shape-inference callback registered for (domain, op_type), or nullptr if none is registered. domain == "" is normalized to :cpp:var:kOnnxDomain.

inline bool RemoveCustomShapeInferenceFunction(const std::string &domain, const std::string &op_type)#: Removes the custom shape-inference callback registered for (domain, op_type). domain == "" is normalized to :cpp:var:kOnnxDomain. Returns true when an entry was removed and false when no callback matched that key.

inline void ClearCustomShapeInferenceFunctions()#: Removes every custom shape-inference callback.

inline const CustomShapeInferenceMap &CustomShapeInferenceFunctions() const noexcept#: Read-only access to all registered custom shape-inference callbacks.

inline bool AddConstraint(const std::string &a, const std::string &b)#: Records that two symbolic dimension names are equal. The pair is canonicalised so that (a, b) and (b, a) are stored only once, and the trivial self-equality is dropped. Returns true when a new constraint was inserted, false otherwise (either a duplicate or a self-constraint).

inline bool HasConstraint(const std::string &a, const std::string &b) const#: true when an equality constraint between a and b is recorded (canonical order is applied before lookup).

inline std::size_t ConstraintsSize() const noexcept#: Number of recorded constraints.

inline const std::set<Constraint> &Constraints() const noexcept#: Read-only access to the underlying set of equality constraints. Each element is a (lhs, rhs) pair with lhs < rhs.

inline bool AddLessEqualConstraint(const std::string &lhs, const std::string &rhs)#: Records that the symbolic dimension named lhs is less-than-or-equal-to the expression rhs. The trivial self-bound (lhs == rhs) is dropped, and empty operands are rejected (they cannot designate a valid dimension name nor a well-formed bound expression). Returns true when a new constraint was inserted, false otherwise (duplicate, self-bound, or empty operand).

inline bool HasLessEqualConstraint(const std::string &lhs, const std::string &rhs) const#: true when a lhs <= rhs constraint is recorded. lhs == rhs always returns true.

inline std::size_t LessEqualConstraintsSize() const noexcept#: Number of recorded <= constraints.

inline const std::set<LessEqualConstraint> &LessEqualConstraints() const noexcept#: Read-only access to the underlying set of <= constraints. Each element is an ordered (lhs, rhs) pair meaning lhs <= rhs.

void ComputeShapeNode(const NodeProto &node)#: Dispatches a single NodeProto to the matching per-operator ComputeShape* function and stores the resulting output :cpp:class:OptimTensor descriptors in *this.

void CheckInputsAvailable(const NodeProto &node) const#: Throws std::invalid_argument if any non-empty input name declared by node is missing from *this.

void CheckOutputsNotAvailable(const NodeProto &node) const#: Throws std::invalid_argument if any non-empty output name declared by node already has an entry in *this.

void ComputeShapes(const utils::RepeatedProtoField<NodeProto> &nodes)#: Runs :cpp:func:ComputeShapeNode on every node of nodes in order.

void ComputeShapeGraph(const GraphProto &graph)#: Seeds *this from the initializers and inputs of graph and then runs :cpp:func:ComputeShapes on its nodes.

void ComputeShapeModel(const ModelProto &model, bool prefill_with_value_info_output = false)#: Runs shape inference on model.graph(), also recording opset versions and local functions from model.

void ApplyInferredShapesToGraph(GraphProto &graph) const#: Writes the shape and element-type descriptors stored in *this back into graph.

void ApplyInferredShapesToModel(ModelProto &model) const#: Writes the shape and element-type descriptors stored in *this back into model.graph().

inline void set_events_enabled(bool enabled) noexcept#: Enables or disables event logging. When disabled (the default), :cpp:func:Set, :cpp:func:ComputeShapeNode, :cpp:func:AddConstraint and :cpp:func:AddLessEqualConstraint skip all event construction. Call set_events_enabled(true) before running shape inference if descriptor or constraint tracing is required.

inline bool events_enabled() const noexcept#

inline void set_current_subgraph(int64_t node_index, const std::string &attr_name)#: Index of the control-flow node in the parent graph currently being inferred. Set before running subgraph shape inference so that events recorded inside carry :cpp:var:ShapeEvent::subgraph_node_index and :cpp:var:ShapeEvent::subgraph_attr_name. -1 for the top-level graph. Use :cpp:func:set_current_subgraph to update both the index and the attribute name atomically.

inline int64_t current_subgraph_node_index() const noexcept#

inline const std::string &current_subgraph_attr_name() const noexcept#

inline void set_current_node_index(int64_t index) noexcept#: Index of the node currently being processed, used to tag the :cpp:var:ShapeEvent::node_index of descriptors and events recorded during its shape-inference dispatch. Set by :cpp:func:ComputeShapes before each :cpp:func:ComputeShapeNode call, -2 while seeding initializers and -1 while seeding graph inputs (or when no node is being processed).

inline int64_t current_node_index() const noexcept#

inline const ShapeEventLog &Events() const noexcept#: Append-only log of every tensor descriptor mutation performed through :cpp:func:Set, every node dispatched through :cpp:func:ComputeShapeNode and every constraint recorded through :cpp:func:AddConstraint / :cpp:func:AddLessEqualConstraint. See :cpp:class:ShapeEvent for the captured fields.

inline ShapeEventLog &Events() noexcept#

inline void ClearEvents() noexcept#: Empties the event log without otherwise touching the context.

inline void set_event_callback(std::function<void(const ShapeEvent&)> callback)#: Registers a callback invoked synchronously for every newly appended :cpp:class:ShapeEvent. Used by callers that want to stream shape inference progress while computation is still running.

inline void clear_event_callback() noexcept#

inline bool has_event_callback() const noexcept#

void AppendComputeNodeEvent(const std::string &op_domain, const std::string &op_type, std::vector<std::string> inputs)#: Appends a :cpp:class:ShapeEvent with action :cpp:enumerator:ShapeEventAction::kComputeNode summarising the shape-inference dispatch of a single NodeProto (its op_domain / op_type and the inputs it consumed). Appended by :cpp:func:ComputeShapeNode for every dispatched node when event logging is enabled.

Private Functions

void LogSetEvent(const std::string &name, const OptimTensor &tensor)#: Appends a :cpp:enumerator:ShapeEventAction::kAdd / :cpp:enumerator:ShapeEventAction::kReplace event for name describing tensor (the descriptor about to be stored). Only called by :cpp:func:Set when event logging is enabled.

void LogConstraintEvent(ShapeEventAction action, const std::string &lhs, const std::string &rhs)#: Appends a :cpp:enumerator:ShapeEventAction::kConstraint / :cpp:enumerator:ShapeEventAction::kConstraintMax event recording the two operands of a newly inserted constraint in inputs. Only called by :cpp:func:AddConstraint / :cpp:func:AddLessEqualConstraint when event logging is enabled.

inline void EmitEvent(const ShapeEvent &ev) const#

Private Members

std::unordered_map<std::string, OptimTensor> tensors_#

std::unordered_map<std::string, OptimSequence> sequences_#

std::unordered_map<std::string, int> opsets_#

std::unordered_map<std::string, const FunctionProto*> local_functions_#

CustomShapeInferenceMap custom_shape_inference_#

std::set<Constraint> constraints_#

std::set<LessEqualConstraint> le_constraints_#

std::map<SubgraphContextKey, std::shared_ptr<ShapesContext>> subgraph_contexts_#

ShapeEventLog events_#

std::function<void(const ShapeEvent&)> event_callback_#

bool events_enabled_ = false#

int64_t current_node_index_ = -1#

int64_t current_subgraph_node_index_ = -1#: Index of the control-flow node in the parent graph currently being inferred (see :cpp:func:set_current_subgraph). -1 for the top-level graph.

std::string current_subgraph_attr_name_#: Attribute name of the subgraph currently being inferred (see :cpp:func:set_current_subgraph). Empty for the top-level graph; set to "body", "then_branch", "else_branch", etc. when running shape inference for a control-flow body subgraph.

std::unordered_map<std::string, std::string> topk_k_dims_#: Tracks the mapping from TopK K-input names to the symbolic dimension name assigned by :cpp:func:TopKKDimName. Populated lazily.

Private Static Functions

static inline std::string NormaliseDomain(const std::string &domain)#