yobx.xshape#

modules

BasicShapeBuilder#

class yobx.xshape.BasicShapeBuilder(verbose: int = 0, opset: int | None = None)[source]#

Implements a basic class doing shape inference in an ONNX model.

A couple of environment variables can be set to help debugging any issue.

ONNXSTOPSHAPE=<name>: raises an exception when name receives a shape.
ONNXSTOPTYPE=<name>: raises an exception when name receives a type.
ONNXDYNDIM=<name>: raises an exception when dimension name is used
ONNXCST=1: shows which constant is requested
ONNXSHAPECOMPUTE=1: raises an exception when a shape is missing
ONNXSTOPVALUESHAPE=<name>: more information in function dealing with shapes

add_dynamic_dimension(name: str)[source]#: Adds a dynamic dimension.

estimate_node_flops(node: NodeProto) → int | str | None[source]#

Estimates the number of floating-point operations for node using the shapes already inferred by run_model().

Uses estimate_node_flops() internally, providing this builder’s get_shape() and get_constant() implementations as the shape_fn and literal_fn callbacks.

Parameters:: node – ONNX node to estimate
Returns:: estimated FLOPs count, or None when shapes are unknown or the op_type is not supported

evaluate_cost_with_true_inputs(feeds: Dict[str, ndarray], cost: List[Tuple[str, int | str | None, Tuple]], exc: bool = False) → List[Tuple[str, int | None, Tuple]][source]#

Evaluates symbolic FLOPs expressions in cost using actual tensor shapes from feeds.

When run_model() is called with InferenceMode.COST on a model that has symbolic (dynamic) input shapes, the returned FLOPs values may be symbolic expressions (strings such as "(DIM1)*(DIM2)*(3)"). This method substitutes the true dimension values extracted from feeds to produce concrete integer FLOPs.

Parameters:

feeds – mapping {name: array} of actual input tensors
cost – list of (op_type, flops, input_shapes) tuples as returned by run_model(..., inference=InferenceMode.COST)
exc – if True, re-raise any evaluation error; otherwise the FLOPs entry is set to None for that node

Returns:

list of (op_type, evaluated_flops, input_shapes) tuples

get_constant(name: str, exc: bool = True, computed_value: bool = False, as_shape: bool = False, multiple_outputs: bool = False) → ndarray | NodeProto[source]#

The method returns the constant name. It is a tensor (numpy array) or a NodeProto which must be evaluated. If computed_value is True, the NodeProto is evaluated with the ReferenceEvaluator.

Parameters:

name – constant name
exc – raise an exception if anything is impossible to do
computed_value – compute the value if not a constant
as_shape – returns a tuple for a shape
multiple_outputs – allow multiple outputs

Returns:

value

get_debug_msg(limit: int = 1000) → str[source]#

Returns a string providing as much information as possible to help the developer understand why a conversion failed.

Parameters:: limit – limit the string if the model is big
Returns:: many pieces of information about the on going conversion

get_device(name) → int[source]#: Returns the device of a result.

get_local_function(name: str, domain: str = '', builder: bool = False) → FunctionProto | BasicShapeBuilder[source]#: Returns a local function.

get_opset(name: str) → int[source]#

Returns the opset version for domain name.

Parameters:: name – domain name
Returns:: domain version or 0 if not specified

get_rank(name: str) → int[source]#: Returns the rank of a result.

get_registered_constraints()[source]#: Returns the constraints registered so far.

get_shape(name: str) → Tuple[int | torch.SymInt | torch.SymFloat | TracingInt | float | str, ...][source]#: Returns the shape of a result.

get_type(name: str) → int[source]#: Returns the type of a result.

has_device(name) → bool[source]#: Tells if a result has a device.

has_local_function(name: str, domain: str = '', builder: bool = False) → bool[source]#: Checks if a local function exists.

has_opset(name: str) → bool[source]#: Tells if opset name is defined.

has_rank(name: str) → bool[source]#: Tells if a result has a rank.

has_shape(name: str, full=False) → bool[source]#: Tells if a result has a shape. If full is True, it returns True if the shape exists and if it is a static shape with all dimensions > 0.

has_type(name: str) → bool | int[source]#: Tells if a result has a type. This should be always true.

property input_names: List[str]#: Returns the list of input names of the model.

is_constant(name: str) → bool[source]#: Tells if a result is a constant.

property output_names: List[str]#: Returns the list of output names of the model.

register_dynamic_objects_from_dim(dim: str)[source]#: Registers all the dynamic objects required in a dimension.

run_model(model: ModelProto | GraphProto, functions: Dict[Tuple[str, str], FunctionProto] | None = None, exc: bool = False, inference: InferenceMode | str = InferenceMode.SHAPE)[source]#

Runs inference over a model or a graph.

Parameters:

model – an ONNX model or graph
functions – a dictionary of functions available to the model
exc – if True, raises an exception when inference fails
inference – InferenceMode value (or its string name, case-insensitive). InferenceMode.SHAPE (default) runs the full shape and type inference using symbolic expressions; InferenceMode.TYPE runs a lighter type-only inference via type_inference.infer_types that only propagates element types without computing shapes

run_node(node: NodeProto, exc: bool = False, cost: bool = True)[source]#: Uses shapes availables in the ShapeBuilder to infer the output shapes and types.

run_value_info(info: ValueInfoProto, is_input: bool)[source]#: Fills ShapeBuilder with information coming from an input or output.

set_constant(name: str, value: TensorProto | NodeProto) → None[source]#: Stores a constant (a onnx.TensorProto or a onnx.NodeProto).

set_device(name: str, device: int | torch.dtype, exc: bool = True)[source]#

Sets the shape for a result. It is exists, it checks the new shape is equal to the existing one.

Parameters:

name – name
device – an integer or a torch device then converted into an integer
exc – raises an exception

set_opset(name: str, version: int)[source]#

Sets the opset version for domain name.

Parameters:

name – domain name
version – domain version

set_rank(name: str, value: int) → bool[source]#

Sets the rank for a result.

Parameters:

name – result name
value – rank

Returns:

True if there is no rank conflict

Sets the shape for a result. It is exists, it checks the new shape is equal to the existing one.

Parameters:

name – result name
shape – shape
exc – raise an exception if inconsistency

set_type(name: str, dtype: int, exc: bool = True) → bool[source]#

Sets the shape for a result. It is exists, it checks the new shape is equal to the existing one.

Parameters:

name – name
dtype – element type (an integer, ONNX), 0 (unknown is a possible value)
exc – raises an exception

Returns:

returns True if there is no type conflict

set_value_shape(name: str, value: Any, equal_to: Tuple[str, str] | None = None)[source]#

Sets the value for a shape result.

Parameters:

name – name
value – it cannot be empty
equal_to – if specified, the value is also equal to this value

A value can be a string (for an unknown shape, a tuple for a shape, an integer for a single scalar.

unique_dimension_name(base: str) → str[source]#: Returns a unique dimension name based on base.

value_as_shape(name: str) → bool[source]#: Returns the value of a result if it is a shape.

ShapeBuilder#

class yobx.xshape.ShapeBuilder[source]#

API for a class computing shapes in an ONNX model.

The main implementation is BasicShapeBuilder. It walks through all the nodes of an ONNX model and infers output shapes and types based on the input shapes, using symbolic expressions when the exact integer values are not known.

Symbolic expressions — When a dimension cannot be determined as a plain integer (e.g. because it depends on a dynamic input dimension), it is stored as a Python-arithmetic string expression built from the input dimension names. For instance, concatenating tensors of shapes ("batch", "seq1") and ("batch", "seq2") along axis 1 yields output shape ("batch", "seq1+seq2"). The supported operators inside a symbolic expression are +, -, *, //, % and ^ (where ^ means max). Expressions are automatically simplified by simplify_expression before being stored, so d + f - f becomes d and 2*seq//2 becomes seq. Once concrete values are available they can be resolved with evaluate_shape() or evaluate_expression.

<<<

import onnx
import onnx.helper as oh
from yobx.xshape import BasicShapeBuilder

TFLOAT = onnx.TensorProto.FLOAT

# Build a small model: Z = Concat(X, Y, axis=1)
model = oh.make_model(
    oh.make_graph(
        [oh.make_node("Concat", ["X", "Y"], ["Z"], axis=1)],
        "graph",
        [
            oh.make_tensor_value_info("X", TFLOAT, ["batch", "seq1"]),
            oh.make_tensor_value_info("Y", TFLOAT, ["batch", "seq2"]),
        ],
        [oh.make_tensor_value_info("Z", TFLOAT, [None, None])],
    ),
    opset_imports=[oh.make_opsetid("", 18)],
    ir_version=10,
)

builder = BasicShapeBuilder()
builder.run_model(model)

print("input names :", builder.input_names)
print("output names:", builder.output_names)
print("shape of Z  :", builder.get_shape("Z"))
print("type of Z   :", builder.get_type("Z"))

>>>

    input names : ['X', 'Y']
    output names: ['Z']
    shape of Z  : ('batch', 'seq1+seq2')
    type of Z   : 1

Constraint mechanism — When a broadcasting operation aligns a symbolic dimension (e.g. "d_model") with a concrete integer (e.g. 64), the concrete value is used immediately as the output dimension and the equality "d_model" = 64 is stored as a constraint. This avoids the need to backtrack through earlier nodes when the concrete value is later discovered. Constraints are inspected with get_registered_constraints() and are used internally for dimension renaming and equality checks. See ShapeBuilder for details.

add_to_constraints(dim_name: str, value: str | int | Set[str | int])[source]#

Adds a constraint associating a symbolic dimension name with a value or set of values.

Parameters:

dim_name – symbolic dimension name (e.g. "batch")
value – the value, name, or set of values/names to associate with that dimension

compare_with_true_inputs(inputs: Dict[str, ndarray] | List[ndarray], outputs: Dict[str, ndarray] | List[ndarray], exc: bool = True, do_shape: bool = True, do_type: bool = True) → Dict[str, Tuple[Tuple[str, int, int], ...]][source]#

Compares the shape of the outputs with what the output shapes would return.

Parameters:

inputs – inputs
outputs – outputs
exc – raises an exception when a discrepancy is met
do_type – compares types
do_shape – compares shapes

Returns:

list of expression, expected value, computed value

evaluate_dimension_equality_with_constraints(d1: str, *args) → bool[source]#: Tells if two dimensions are equal.

get_attribute(node: NodeProto, att_name: str, exc: bool = True) → AttributeProto | None[source]#: Returns an attribute for a node.

get_attribute_with_default(node: NodeProto, name: str, default_value: Any) → Any[source]#

Returns an attribute or its default value if missing.

Parameters:

node – node
name – attribute name
default_value – default value

Returns:

value

get_attributes_with_default(node: NodeProto, **default_values) → Dict[str, Any][source]#

Returns int or float attributes. If missing, the default value is returned if it is not None.

Parameters:

node – node
default_values – default values

get_debug_msg(limit: int = 1000) → str[source]#

Returns a string providing as much information as possible to help the developer understand why a conversion failed.

Parameters:: limit – limit the string if the model is big
Returns:: many pieces of information about the on going conversion

get_device(name: str) → int[source]#

Returns the device of result name.

Parameters:: name – result name
Returns:: rank as an integer

get_opset(name: str) → int[source]#

Returns the opset version for domain name.

Parameters:: name – domain name
Returns:: domain version or 0 if not defined

get_rank(name: str) → int[source]#

Returns the rank (number of dimensions) of result name.

Parameters:: name – result name
Returns:: rank as an integer

get_registered_constraints() → Dict[str, Set[str | int]][source]#

Returns the constraints registered so far.

Returns:: mapping from dimension name to the set of values/names it is constrained to be equal to

Returns the shape of result name as a tuple. Each dimension is either an integer or a string (symbolic dimension).

Parameters:: name – result name
Returns:: shape as a tuple of integers and/or strings

Returns the shape of result name using user-visible dimension names.

After _improves_dynamic_dimension_naming() has been called, symbolic dimension names that were given by the user (e.g. "batch", "seq_length") are substituted for the internal names (e.g. "s0", "s1"). When no renaming has been computed yet this falls back to get_shape().

Parameters:: name – result name
Returns:: shape tuple with user dimension names where available

get_type(name: str) → int[source]#

Returns the element type of result name as an ONNX integer (e.g. onnx.TensorProto.FLOAT == 1).

Parameters:: name – result name
Returns:: element type as an integer

has_device(name: str) → bool[source]#: Tells if name has a device.

has_opset(name: str) → bool[source]#: Tells if opset name is defined.

has_rank(name: str) → bool[source]#: Tells if name has a rank.

has_shape(name: str) → bool[source]#: Tells if name has a shape.

has_type(name: str) → bool[source]#: Tells if name has a type.

property input_names: List[str]#: Returns the list of input names of the model.

property output_names: List[str]#: Returns the list of output names of the model.

pretty_node(node: NodeProto | None, limit: int = 80, short: bool = True, shape: bool = False) → str[source]#

Pretty rendering for a node.

Parameters:

node – node to render
limit – to show type and shapes after the limit
short – do not display shape information on the left
shape – show shape information below

Returns:

string

register_constraint_dimension(dim_name: str, value: Any)[source]#

Registers a constraint associating a symbolic dimension name with a value. This allows to deal backward constraints after a single pass if the model.

Parameters:

dim_name – symbolic dimension name (e.g. "batch")
value – the value or set of values to associate with that dimension

set_device(name: str, rank: int)[source]#

Sets the rank (number of dimensions) for result name.

Parameters:

name – result name
rank – rank as an integer

set_opset(name: str, itype: int)[source]#

Sets the opset version for domain name.

Parameters:

name – domain name
version – domain version

set_rank(name: str, rank: int)[source]#

Sets the rank (number of dimensions) for result name.

Parameters:

name – result name
rank – rank as an integer

Sets the shape for result name.

Parameters:

name – result name
shape – tuple of integers and/or strings (symbolic dimensions)

set_type(name: str, itype: int)[source]#

Sets the element type for result name.

Parameters:

name – result name
itype – element type as an ONNX integer (e.g. onnx.TensorProto.FLOAT == 1)

update_shapes(model: ModelProto)[source]#: Updates model shapes with the value stored inside this graph.

estimate_node_flops#

yobx.xshape.estimate_node_flops(node: NodeProto, shape_fn: Callable[[str], Tuple[int | str, ...] | None], literal_fn: Callable[[str], Tuple[int | str, ...] | None]) → int | str | None[source]#

Estimates the number of floating-point operations for a single ONNX node.

Returns None when the shapes are not fully known (dynamic shapes) or the op_type is not covered.

Parameters:

node – ONNX node
shape_fn – callable mapping tensor name → shape tuple (from shape inference)
literal_fn – callable mapping tensor name → int-value tuple for 1-D integer constant tensors (shape specification tensors); used as a fallback when shape_fn cannot resolve a shape

Returns:

estimated number of FLOPs, or None

list_op_cost_formulas#

yobx.xshape.list_op_cost_formulas() → Dict[str, str][source]#

Returns a mapping from each supported ONNX op_type to the symbolic FLOPs expression produced by estimate_node_flops() on a representative test case from the ONNX backend test suite.

For every single-node model found in the ONNX backend test data directory the static input dimensions are replaced by symbolic variables (DIM<n>) using replace_static_dimensions_by_strings(). BasicShapeBuilder is then run with inference=InferenceMode.COST to obtain the symbolic FLOPs expression.

Only the first passing test case per op_type is kept. Operators with no matching backend test case, or whose cost cannot be inferred symbolically, are omitted.

Returns:: {op_type: symbolic_flops_expression} sorted alphabetically.

yobx.xshape#

BasicShapeBuilder#

ShapeBuilder#

estimate_node_flops#

list_op_cost_formulas#

This Page