onnx_diagnostic.torch_export_patches.patch_expressions

onnx_diagnostic.torch_export_patches.patch_expressions.make_undefined_dimension(i: int) → SymInt

Uses for a custom op when a new dimension must be introduced to bypass some verification. The following function creates a dummy output with a dimension based on the content.

def symbolic_shape(x, y):
    return torch.empty(
        x.shape[0],
        make_undefined_dimension(min(x.shape[1], y[0])),
    )

onnx_diagnostic.torch_export_patches.patch_expressions.patched_float_arange(start, end, step)

Patched arange when start, end, step are floats.

onnx_diagnostic.torch_export_patches.patch_expressions.patched_selector(fct: Callable, patched_fct: Callable) → Callable

Returns fct if the model is being executed or patched_fct if it is being exported.

onnx_diagnostic.torch_export_patches.patch_expressions.register_patched_expressions(namespace: str = 'patched')

Registers as custom ops known expressions failing due to dynamic shapes.

<<<

import pprint
from onnx_diagnostic.torch_export_patches.patch_expressions import (
    _iterate_patched_expressions,
)

pprint.pprint([name for name, _f, _fsh in _iterate_patched_expressions()])

>>>

    ['float_arange']