yobx.xshape¶

modules

BasicShapeBuilder¶

class yobx.xshape.BasicShapeBuilder(verbose: int = 0, opset: int | None = None)[source][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][source]¶: Adds a dynamic dimension.

get_constant(name: str, exc: bool = True, computed_value: bool = False, as_shape: bool = False, multiple_outputs: bool = False) → ndarray | NodeProto[source][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][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_rank(name: str) → int[source][source]¶: Returns the rank of a result.

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

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

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

has_shape(name: str, full=False) → bool[source][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][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][source]¶: Tells if a result is a constant.

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

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

Registers a constraint on a dimension.

Parameters:

dim_name – dimension name
value – value to register

register_dynamic_objects_from_dim(dim: str)[source][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)[source][source]¶: Runs inference over a model or a graph.

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

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

set_constant(name: str, value: TensorProto | NodeProto) → None[source][source]¶: Tells if a result is a constant.

set_device(name: str, device: int | torch.dtype, exc: bool = True)[source][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_rank(name: str, value: int) → bool[source][source]¶

Sets the rank for a result.

Parameters:

name – result name
value – rank

Returns:

True if there is no rank conflict

set_shape(name: str, shape: Tuple[int | torch.SymInt | torch.SymFloat | float | str, ...], exc: bool = False, **_kwargs)[source][source]¶

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][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][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][source]¶: Returns a unique dimension name based on base.

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

ShapeBuilder¶

class yobx.xshape.ShapeBuilder[source][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

compare_with_true_inputs(inputs: Dict[str, ndarray] | List[ndarray], outputs: Dict[str, ndarray] | List[ndarray], exc: bool = True) → Dict[str, Tuple[Tuple[str, int, int], ...]][source][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

Returns:

list of expression, expected value, computed value

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

get_attribute_with_default(node: NodeProto, name: str, default_value: Any) → Any[source][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][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_rank(name: str) → int[source][source]¶

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

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

get_shape(name: str) → Tuple[int | torch.SymInt | torch.SymFloat | float | str, ...][source][source]¶

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

get_type(name: str) → int[source][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

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][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][source]¶

Registers a constraint associating a symbolic dimension name with a value.

Parameters:

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

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

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

Parameters:

name – result name
rank – rank as an integer

set_shape(name: str, shape: Tuple[int | torch.SymInt | torch.SymFloat | float | str, ...])[source][source]¶

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][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][source]¶: Updates model shapes with the value stored inside this graph.