Source code for experimental_experiment.torch_interpreter

from ._exceptions import FunctionNotFoundError
from .export_options import ExportOptions
from .onnx_export import to_onnx, match_input_parameters, FunctionOptions
from .dispatcher import Dispatcher, ForceDispatcher

LOCAL_DOMAIN = "aten_local_function"



[docs]
class TorchOpOverload:
    """
    The class is unused only to bypass a documentation warning.
    The alias ``TorchOpOverload`` refers to ``torch._ops.Overload``.
    """

    pass  # noqa: PIE790




[docs]
def make_undefined_dimension(i: int) -> "torch.SymInt":  # noqa: F821
    """
    Uses for a custom op when a new dimension must be introduced to bypass
    some verficiation. The following function creates a dummy output
    with a dimension based on the content.

    .. code-block:: python

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

    """
    import torch

    t = torch.ones((i * 2,))
    t[:i] = 0
    res = torch.nonzero(t).shape[0]
    return res