onnx_diagnostic.export¶

modules

CoupleInputsDynamicShapes¶

class onnx_diagnostic.export.CoupleInputsDynamicShapes(args: Tuple[Any, ...], kwargs: Dict[str, Any], dynamic_shapes: Tuple[Tuple[Any, ...], Dict[str, Any]], args_names: Module | List[str] | None = None)[source]¶

Pair inputs / dynamic shapes.

Parameters:

args – positional arguments
kwargs – named arguments
dynamic_shapes – dynamic shapes
args_names – if both args and kwargs are not empty, then dynamic shapes must be a dictionary, and positional must be added to the named arguments. Arguments names or a module must be given in that case.

change_dynamic_dimensions(desired_values: Dict[str, int] | None = None, args_kwargs: bool = False)[source]¶

A model exported with dynamic shapes is not necessarily dynamic just because the user specified dynamic shapes. The algorithm may discover that a dimension cannot be dynamic and then continues the export making the assumption it is static. That may lead a wrong model. This function produces a new set of inputs with different values for the dimension than the first ones, assuming they were used to export the model.

Parameters:

desired_values – to fixed named dimension to have the desired value
args_kwargs – return both args, kwargs even if empty

Returns:

new inputs

Example:

<<<

import torch
from onnx_diagnostic.helpers import string_type
from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes

T3x15 = torch.rand((3, 15))
T3x20 = torch.rand((3, 20))
T3x4 = torch.rand((3, 4))
ds_batch = {0: "batch"}
ds_batch_seq = {0: "batch", 1: "seq"}
kwargs = {"A": T3x4, "B": (T3x15, T3x20)}
ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
new_kwargs = CoupleInputsDynamicShapes((), kwargs, ds).change_dynamic_dimensions()
print("before:", string_type(kwargs, with_shape=True))
print("-after:", string_type(new_kwargs, with_shape=True))

>>>

    before: dict(A:T1s3x4,B:(T1s3x15,T1s3x20))
    -after: dict(A:T1s4x4,B:(T1s4x15,T1s4x21))

invalid_dimensions_for_export()[source]¶

Tells if the inputs are valid based on the dynamic shapes definition. The method assumes that all custom classes can be serialized. If some patches were applied to export, they should enabled while calling this method if the inputs contains such classes.

The function checks that a dynamic dimension does not receive a value of 0 or 1. It returns the unexpected values in the same structure as the given dynamic shapes.

Example:

<<<

import torch
from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes

T3x1 = torch.rand((3, 1))
T3x4 = torch.rand((3, 4))
ds_batch = {0: "batch"}
ds_batch_seq = {0: "batch", 1: "seq"}
kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
print(CoupleInputsDynamicShapes((), kwargs, ds).invalid_dimensions_for_export())

>>>

    {'B': (None, {1: 'd=[1]'})}

In case it works, it shows:

<<<

import torch
from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes

T3x2 = torch.rand((3, 2))
T3x4 = torch.rand((3, 4))
ds_batch = {0: "batch"}
ds_batch_seq = {0: "batch", 1: "seq"}
kwargs = {"A": T3x4, "B": (T3x2, T3x2)}
ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
print(CoupleInputsDynamicShapes((), kwargs, ds).invalid_dimensions_for_export())

>>>

    None

replace_by_string()[source]¶

Replaces dimensions by strings.

Example:

<<<

import torch
from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes

Dim = torch.export.Dim
T3x1 = torch.rand((3, 1))
T3x4 = torch.rand((3, 4))
ds_batch = {0: Dim("batch")}
ds_batch_seq = {0: Dim("batch"), 1: Dim("seq")}
kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
print(CoupleInputsDynamicShapes((), kwargs, ds).replace_by_string())

>>>

    {'A': {0: 'batch'}, 'B': ({0: 'batch'}, {0: 'batch', 1: 'seq'})}

replace_string_by(value: Any = None)[source]¶

Replaces string by the value torch.export.Dim.DYNAMIC (default) or any other value specified by value.

Example:

<<<

import torch
from onnx_diagnostic.export.dynamic_shapes import CoupleInputsDynamicShapes

T3x1 = torch.rand((3, 1))
T3x4 = torch.rand((3, 4))
ds_batch = {0: "batch"}
ds_batch_seq = {0: "batch", 1: "seq"}
kwargs = {"A": T3x4, "B": (T3x1, T3x1)}
ds = {"A": ds_batch, "B": (ds_batch, ds_batch_seq)}
print(CoupleInputsDynamicShapes((), kwargs, ds).replace_string_by())

>>>

    {'A': {0: _DimHint(type=<_DimHintType.DYNAMIC: 3>, min=None, max=None, _factory=True)}, 'B': ({0: _DimHint(type=<_DimHintType.DYNAMIC: 3>, min=None, max=None, _factory=True)}, {0: _DimHint(type=<_DimHintType.DYNAMIC: 3>, min=None, max=None, _factory=True), 1: _DimHint(type=<_DimHintType.DYNAMIC: 3>, min=None, max=None, _factory=True)})}

ModelInputs¶

class onnx_diagnostic.export.ModelInputs(model: Module, inputs: List[Tuple[Any, ...]] | List[Dict[str, Any]] | List[Tuple[Tuple[Any, ...], Dict[str, Any]]], level: int = 0, method_name: str = 'forward', name: str = 'main')[source]¶

Wraps a model and a couple of sets of valid inputs. Based on that information, the class is able to infer the dynamic shapes for torch.export.export().

Parameters:

model – model to export
inputs – list of valid set of inputs
level – if this module is a submodule, it is the level of submodule
method_name – by default, the forward method is processed but it could be another one
name – a name, mostly for debugging purposes

Examples:

args

<<<

import pprint
import torch
from onnx_diagnostic.export import ModelInputs


class Model(torch.nn.Module):
    def forward(self, x, y):
        return x + y


model = Model()
x = torch.randn((5, 6))
y = torch.randn((1, 6))
model(x, y)  # to check it works

inputs = [(x, y), (torch.randn((7, 8)), torch.randn((1, 8)))]
mi = ModelInputs(Model(), inputs)
ds = mi.guess_dynamic_shapes()
pprint.pprint(ds)

>>>

    (({0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                   min=None,
                   max=None,
                   _factory=True),
       1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                   min=None,
                   max=None,
                   _factory=True)},
      {1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                   min=None,
                   max=None,
                   _factory=True)}),
     {})

kwargs

<<<

import pprint
import torch
from onnx_diagnostic.export import ModelInputs


class Model(torch.nn.Module):
    def forward(self, x, y):
        return x + y


model = Model()
x = torch.randn((5, 6))
y = torch.randn((1, 6))
model(x=x, y=y)  # to check it works

inputs = [dict(x=x, y=y), dict(x=torch.randn((7, 8)), y=torch.randn((1, 8)))]
mi = ModelInputs(Model(), inputs)
ds = mi.guess_dynamic_shapes()
pprint.pprint(ds)

>>>

    ((),
     {'x': {0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True),
            1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True)},
      'y': {1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True)}})

args and kwargs

<<<

import pprint
import torch
from onnx_diagnostic.export import ModelInputs


class Model(torch.nn.Module):
    def forward(self, x, y):
        return x + y


model = Model()
x = torch.randn((5, 6))
y = torch.randn((1, 6))
model(x, y=y)  # to check it works

inputs = [((x,), dict(y=y)), ((torch.randn((7, 8)),), dict(y=torch.randn((1, 8))))]
mi = ModelInputs(Model(), inputs)
ds = mi.guess_dynamic_shapes()
pprint.pprint(ds)

>>>

    (({0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                   min=None,
                   max=None,
                   _factory=True),
       1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                   min=None,
                   max=None,
                   _factory=True)},),
     {'y': {1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True)}})

torch.export.export() does not like dynamic shapes defined both as args and kwargs. kwargs must be used. move_to_kwargs modifies the inputs and the dynamic shapes to make the model and the given inputs exportable.

<<<

import pprint
import torch
from onnx_diagnostic.export import ModelInputs
from onnx_diagnostic.helpers import string_type


class Model(torch.nn.Module):
    def forward(self, x, y):
        return x + y


model = Model()
x = torch.randn((5, 6))
y = torch.randn((1, 6))
model(x, y=y)  # to check it works

inputs = [((x,), dict(y=y)), ((torch.randn((7, 8)),), dict(y=torch.randn((1, 8))))]
mi = ModelInputs(Model(), inputs)
ds = mi.guess_dynamic_shapes()

a, kw, nds = mi.move_to_kwargs(*mi.inputs[0], ds)
print("moved args:", string_type(a, with_shape=True))
print("moved kwargs:", string_type(kw, with_shape=True))
print("dynamic shapes:")
pprint.pprint(nds)

>>>

    moved args: (T1s5x6,)
    moved kwargs: dict(y:T1s1x6)
    dynamic shapes:
    ((),
     {'x': {0: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True),
            1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True)},
      'y': {1: _DimHint(type=<_DimHintType.DYNAMIC: 3>,
                        min=None,
                        max=None,
                        _factory=True)}})

property full_name: str¶: Returns a name and class name.

guess_dynamic_dimensions(*tensors, auto: bool = False) → Dict[int, Any] | None[source]¶: Infers the dynamic dimension from multiple shapes. If auto is True, it returns torch.export.Dim.AUTO for every dimension which cannot be guessed. Two tensors with the same value for one dimension can be guessed, but if there is only 1, it cannot.

guess_dynamic_shape_object(*objs: Any, auto: bool = False, msg: Callable | None = None) → Any[source]¶: Guesses the dynamic shapes for one argument.

guess_dynamic_shapes(auto: bool = False) → Tuple[Tuple[Any, ...], Dict[str, Any]][source]¶

Guesses the dynamic shapes for that module from two execution. If there is only one execution, then that would be static dimensions.

Parameters:: auto – if auto is True, use torch.export.Dim.AUTO for any dimension if the number of inputs is one

property module_name_type¶: Returns name and module type.

move_to_kwargs(args: Tuple[Any, ...], kwargs: Dict[str, Any], dynamic_shapes: Tuple[Tuple[Any, ...], Dict[str, Any]]) → Tuple[Tuple[Any, ...], Dict[str, Any], Tuple[Tuple[Any, ...], Dict[str, Any]]][source]¶: Uses the signatures to move positional arguments (args) to named arguments (kwargs) with the corresponding dynamic shapes. kwargs, dynamic_shapes are modified inplace.

process_inputs(inputs: List[Tuple[Any, ...]] | List[Dict[str, Any]] | List[Tuple[Tuple[Any, ...], Dict[str, Any]]]) → List[Tuple[Tuple[Any, ...], Dict[str, Any]]][source]¶: Transforms a list of valid inputs, list of args, list of kwargs or list of both into a list of (args, kwargs).

property true_model_name: str¶: Returns class name or module name.

validate_inputs_for_export(dynamic_shapes: Tuple[Tuple[Any, ...], Dict[str, Any]] | None = None) → List[List[int | str]][source]¶

Validates the inputs the class contains for the given dynamic shapes. If not specified, the dynamic_shapes are guessed.

Parameters:: dynamic_shapes – dynamic shapes to validate
Returns:: a list of lists, every list contains the path the invalid dimension

validate_ep¶

onnx_diagnostic.export.validate_ep(ep: Module | ExportedProgram, mod: Module | None = None, args: Tuple[Any, ...] | None = None, kwargs: Dict[str, Any] | None = None, copy: bool = False, dynamic_shapes: Any | None = None, values_to_try: Dict[str, List[int]] | None = None, exc: bool = True, verbose: int = 0, atol: float = 0.01, rtol: float = 0.1) → List[Dict[str, Any]][source]¶

Validates an exported program.

Parameters:

model – first module
mod – second module (it produces the expected values)
args – positional arguments
kwargs – named arguments
copy – copy the inputs before executing the model (they may modify them inplace)
dynamic_shapes – dynamic shapes, string should be used not torch.export.Dim
values_to_try – dictionary with the values to try for every dynamic dimension
exc – raise exception if discrepancies are too high
verbose – verbosity level
atol – absolute tolerance
rtol – relative tolerance

Returns:

dictionary with inputs, outputs and tolerance

onnx_diagnostic.export¶

CoupleInputsDynamicShapes¶

ModelInputs¶

validate_ep¶

Other functions¶