Table of Contents
Shortcuts

# Source code for torch_tensorrt._Input

    from __future__ import annotations

    from enum import Enum
    from typing import Any, Dict, List, Optional, Sequence, Tuple

    import torch
    from torch_tensorrt._enums import dtype, memory_format

    class Input(object):
        """
        Defines an input to a module in terms of expected shape, data type and tensor format.

        Attributes:
            shape_mode (torch_tensorrt.Input._ShapeMode): Is input statically or dynamically shaped
            shape (Tuple or Dict): Either a single Tuple or a dict of tuples defining the input shape.
                Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form

                .. code-block:: py

                    {"min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple}

            dtype (torch_tensorrt.dtype): The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
            format (torch_tensorrt.TensorFormat): The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
        """

        class _ShapeMode(Enum):
            STATIC = 0
            DYNAMIC = 1

        shape_mode: Optional[_ShapeMode] = (
            None  #: Is input statically or dynamically shaped
        )
        shape: Optional[Tuple[int, ...] | Dict[str, Tuple[int, ...]]] = (
            None  #: Either a single Tuple or a dict of tuples defining the input shape. Static shaped inputs will have a single tuple. Dynamic inputs will have a dict of the form ``{ "min_shape": Tuple, "opt_shape": Tuple, "max_shape": Tuple }``
        )
        dtype: dtype = (
            dtype.unknown
        )  #: The expected data type of the input tensor (default: torch_tensorrt.dtype.float32)
        _explicit_set_dtype: bool = False
        format: memory_format = (
            memory_format.linear
        )  #: The expected format of the input tensor (default: torch_tensorrt.memory_format.linear)

        DOMAIN_OFFSET: float = 2.0
        low_tensor_domain_incl: float = 0.0
        high_tensor_domain_excl: float = low_tensor_domain_incl + DOMAIN_OFFSET
        torch_tensor: torch.Tensor = None
        name: str = ""
        is_shape_tensor: bool = False

        def __init__(self, *args: Any, **kwargs: Any) -> None:
            """__init__ Method for torch_tensorrt.Input

            Input accepts one of a few construction patterns

            Args:
                shape (Tuple or List, optional): Static shape of input tensor

            Keyword Arguments:
                shape (Tuple or List, optional): Static shape of input tensor
                min_shape (Tuple or List, optional): Min size of input tensor's shape range
                    Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implicitly this sets Input's shape_mode to DYNAMIC
                opt_shape (Tuple or List, optional): Opt size of input tensor's shape range
                    Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implicitly this sets Input's shape_mode to DYNAMIC
                max_shape (Tuple or List, optional): Max size of input tensor's shape range
                    Note: All three of min_shape, opt_shape, max_shape must be provided, there must be no positional arguments, shape must not be defined and implicitly this sets Input's shape_mode to DYNAMIC
                dtype (torch.dtype or torch_tensorrt.dtype): Expected data type for input tensor (default: torch_tensorrt.dtype.float32)
                format (torch.memory_format or torch_tensorrt.TensorFormat): The expected format of the input tensor (default: torch_tensorrt.TensorFormat.NCHW)
                tensor_domain (Tuple(float, float), optional): The domain of allowed values for the tensor, as interval notation: [tensor_domain[0], tensor_domain[1]).
                    Note: Entering "None" (or not specifying) will set the bound to [0, 2)
                torch_tensor (torch.Tensor): Holds a corresponding torch tensor with this Input.
                name (str, optional): Name of this input in the input nn.Module's forward function. Used to specify dynamic shapes for the corresponding input in dynamo tracer.
            Examples:
                - Input([1,3,32,32], dtype=torch.float32, format=torch.channel_last)
                - Input(shape=(1,3,32,32), dtype=torch_tensorrt.dtype.int32, format=torch_tensorrt.TensorFormat.NCHW)
                - Input(min_shape=(1,3,32,32), opt_shape=[2,3,32,32], max_shape=(3,3,32,32)) #Implicitly dtype=torch_tensorrt.dtype.float32, format=torch_tensorrt.TensorFormat.NCHW
            """
            # Compatibility code for switching over from InputTensorSpec
            if "shape" in kwargs and "shape_ranges" in kwargs:
                assert (
                    len(kwargs["shape_ranges"]) == 1 and len(kwargs["shape_ranges"][0]) == 3
                )
                del kwargs["shape"]

                kwargs["min_shape"] = kwargs["shape_ranges"][0][0]
                kwargs["opt_shape"] = kwargs["shape_ranges"][0][1]
                kwargs["max_shape"] = kwargs["shape_ranges"][0][2]

            if len(args) == 1:
                if not Input._supported_input_size_type(args[0]):
                    raise TypeError(
                        "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
                        + str(type(args[0]))
                    )
                if any(k in kwargs for k in ["min_shape", "opt_shape", "max_shape"]):
                    raise ValueError(
                        "Found that both shape (as a positional argument), and one or more of min_shape, opt_shape, max_shape were specified\nclass Input expects that only either shape or all three of min_shape, opt_shape, max_shape are defined"
                    )
                self.shape = tuple(args[0])
                self.shape_mode = Input._ShapeMode.STATIC

            elif len(args) == 0:
                if "shape" not in kwargs and not (
                    all(k in kwargs for k in ["min_shape", "opt_shape", "max_shape"])
                ):
                    raise ValueError(
                        "Missing required arguments for class Input\nEither shape or all three of min_shape, opt_shape, max_shape must be defined"
                    )
                elif ("shape" in kwargs) and all(
                    k in kwargs for k in ["min_shape", "opt_shape", "max_shape"]
                ):
                    raise ValueError(
                        "Found that both shape, and one or more of min_shape, opt_shape, max_shape were specified\nclass Input expects that only either shape or all three of min_shape, opt_shape, max_shape are defined"
                    )

                if "shape" in kwargs:
                    if not Input._supported_input_size_type(kwargs["shape"]):
                        raise TypeError(
                            "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
                            + str(type(kwargs["shape"]))
                        )
                    self.shape = tuple(kwargs["shape"])
                    self.shape_mode = Input._ShapeMode.STATIC
                else:
                    if not Input._supported_input_size_type(kwargs["min_shape"]):
                        raise TypeError(
                            "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
                            + str(type(kwargs["min_shape"]))
                            + " for min_shape"
                        )
                    if not Input._supported_input_size_type(kwargs["opt_shape"]):
                        raise TypeError(
                            "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
                            + str(type(kwargs["opt_shape"]))
                            + " for opt_shape"
                        )
                    if not Input._supported_input_size_type(kwargs["max_shape"]):
                        raise TypeError(
                            "Input shape specifications for inputs are required to be a List, tuple or torch.Size, found type: "
                            + str(type(kwargs["max_shape"]))
                            + " for max_shape"
                        )

                    self.shape = {
                        "min_shape": tuple(kwargs["min_shape"]),
                        "opt_shape": tuple(kwargs["opt_shape"]),
                        "max_shape": tuple(kwargs["max_shape"]),
                    }
                    self.shape_mode = Input._ShapeMode.DYNAMIC

            else:
                raise ValueError(
                    f"Unexpected number of positional arguments for class Input \n    Found {len(args)} arguments, expected either zero or a single positional arguments"
                )

            if "dtype" in kwargs:
                self.dtype = dtype._from(kwargs["dtype"])

            if self.dtype != dtype.unknown:
                self._explicit_set_dtype = True
            else:
                self._explicit_set_dtype = False

            if "is_shape_tensor" in kwargs:
                self.is_shape_tensor = kwargs["is_shape_tensor"]

            if "format" in kwargs:
                self.format = memory_format._from(kwargs["format"])

            if "tensor_domain" in kwargs:
                domain = kwargs["tensor_domain"]
            else:
                domain = None

            self.tensor_domain = Input._parse_tensor_domain(domain)

            if "torch_tensor" in kwargs:
                self.torch_tensor = kwargs["torch_tensor"]
            else:
                if self.is_shape_tensor:
                    self.torch_tensor = torch.tensor(
                        kwargs["opt_shape"], dtype=kwargs["dtype"]
                    )
                elif self.shape_mode == Input._ShapeMode.DYNAMIC:
                    self.torch_tensor = self.example_tensor("opt_shape")
                else:
                    self.torch_tensor = self.example_tensor()

            if "name" in kwargs:
                self.name = kwargs["name"]

        def __str__(self) -> str:
            if self.shape_mode == Input._ShapeMode.STATIC:
                return "Input(shape={}, dtype={}, format={}, domain=[{}, {}))".format(
                    self.shape,
                    str(self.dtype),
                    str(self.format),
                    str(self.tensor_domain[0]),
                    str(self.tensor_domain[1]),
                )
            elif self.shape_mode == Input._ShapeMode.DYNAMIC:
                if isinstance(self.shape, dict):
                    return "Input(min_shape={}, opt_shape={}, max_shape={}, dtype={}, format={}, domain=[{}, {}))".format(
                        self.shape["min_shape"],
                        self.shape["opt_shape"],
                        self.shape["max_shape"],
                        str(self.dtype),
                        str(self.format),
                        str(self.tensor_domain[0]),
                        str(self.tensor_domain[1]),
                    )
                else:
                    raise RuntimeError(
                        f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})"
                    )
            else:
                raise RuntimeError("Unknown input shape mode")

        def __repr__(self) -> str:
            return self.__str__()

        @staticmethod
        def equivalent_spec(a: Input, b: Input) -> bool:
            if a.shape_mode != b.shape_mode:
                return False

            if a.shape_mode == Input._ShapeMode.DYNAMIC:
                assert isinstance(a.shape, dict)
                assert isinstance(b.shape, dict)
                checks = [
                    a.shape["min_shape"] == b.shape["min_shape"],
                    a.shape["opt_shape"] == b.shape["opt_shape"],
                    a.shape["max_shape"] == b.shape["max_shape"],
                    a.dtype == b.dtype,
                    a.format == b.format,
                    a.low_tensor_domain_incl == b.low_tensor_domain_incl,
                    a.high_tensor_domain_excl == b.high_tensor_domain_excl,
                ]
                return all(checks)
            else:
                checks = [
                    a.shape == b.shape,
                    a.dtype == b.dtype,
                    a.format == b.format,
                    a.low_tensor_domain_incl == b.low_tensor_domain_incl,
                    a.high_tensor_domain_excl == b.high_tensor_domain_excl,
                ]
                return all(checks)

        @staticmethod
        def _supported_input_size_type(input_size: Any) -> bool:
            if isinstance(input_size, torch.Size):
                return True
            elif isinstance(input_size, tuple):
                return True
            elif isinstance(input_size, list):
                return True
            else:
                return False

        @staticmethod
        def _parse_tensor_domain(
            domain: Optional[Tuple[float, float]]
        ) -> Tuple[float, float]:
            """
            Produce a tuple of integers which specifies a tensor domain in the interval format: [lo, hi)

            Args:
                domain (Tuple[int, int]): A tuple of integers (or NoneTypes) to verify

            Returns:
                A tuple of two int32_t-valid integers
            """
            if domain is None:
                result_domain = (
                    Input.low_tensor_domain_incl,
                    Input.high_tensor_domain_excl,
                )

            elif len(domain) == 2:
                domain_lo, domain_hi = domain

                # Validate type and provided values for domain
                valid_type_lo = isinstance(domain_lo, (int, float))
                valid_type_hi = isinstance(domain_hi, (int, float))

                if not valid_type_lo:
                    raise ValueError(
                        f"Expected value for tensor domain low specifier, got {domain_lo}"
                    )
                elif not valid_type_hi:
                    raise ValueError(
                        f"Expected value for tensor domain high specifier, got {domain_hi}"
                    )

                if domain_hi <= domain_lo:
                    raise ValueError(
                        "Expected provided integer range to have low tensor domain value "
                        + f"< high tensor domain value, got invalid range [{domain_lo}, {domain_hi})"
                    )
                result_domain = (float(domain_lo), float(domain_hi))
            else:
                raise ValueError(
                    f"Expected 2 values for domain, got {len(domain)}: {domain}"
                )

            return result_domain

        @classmethod
        def from_tensor(
            cls, t: torch.Tensor, disable_memory_format_check: bool = False
        ) -> "Input":
            """
            Produce a Input which contains the information of the given PyTorch tensor.

            Args:
                tensor (torch.Tensor): A PyTorch tensor.
                disable_memory_format_check (bool): Whether to validate the memory formats of input tensors

            Returns:
                A Input object.
            """
            if not (
                disable_memory_format_check
                or t.is_contiguous(memory_format=torch.contiguous_format)
                or t.is_contiguous(memory_format=torch.channels_last)
            ):
                raise ValueError(
                    "Tensor does not have a supported memory format, supported formats are contiguous or channel_last"
                )
            frmt = (
                torch.contiguous_format
                if (
                    disable_memory_format_check
                    or t.is_contiguous(memory_format=torch.contiguous_format)
                )
                else torch.channels_last
            )
            return cls(shape=t.shape, dtype=t.dtype, format=frmt, torch_tensor=t)

        @classmethod
        def from_tensors(
            cls, ts: Sequence[torch.Tensor], disable_memory_format_check: bool = False
        ) -> List["Input"]:
            """
            Produce a list of Inputs which contain
            the information of all the given PyTorch tensors.

            Args:
                tensors (Iterable[torch.Tensor]): A list of PyTorch tensors.
                disable_memory_format_check (bool): Whether to validate the memory formats of input tensors

            Returns:
                A list of Inputs.
            """

            assert isinstance(ts, (list, tuple))
            return [
                cls.from_tensor(t, disable_memory_format_check=disable_memory_format_check)
                for t in ts
            ]

        def example_tensor(
            self, optimization_profile_field: Optional[str] = None
        ) -> torch.Tensor:
            """
            Get an example tensor of the shape specified by the Input object

            Args:
                optimization_profile_field (Optional(str)): Name of the field to use for shape in the case the Input is dynamically shaped

            Returns:
                A PyTorch Tensor
            """
            if self.shape_mode == Input._ShapeMode.STATIC:
                if optimization_profile_field is not None:
                    raise ValueError(
                        "Specified a optimization profile field but the input is static"
                    )
                else:
                    if isinstance(self.shape, tuple):
                        return torch.rand(self.shape).to(
                            dtype=self.dtype.to(torch.dtype, use_default=True)
                        )
                    else:
                        RuntimeError(
                            f"Input shape is dynamic but shapes are not provided as sequence (found: {self.shape})"
                        )
            else:
                if optimization_profile_field is not None:
                    try:
                        assert any(
                            optimization_profile_field == field_name
                            for field_name in ["min_shape", "opt_shape", "max_shape"]
                        )
                    except AssertionError:
                        raise ValueError(
                            "Invalid field name, expected one of min_shape, opt_shape, max_shape"
                        )

                    if isinstance(self.shape, dict):
                        return torch.rand(self.shape[optimization_profile_field]).to(
                            dtype=self.dtype.to(torch.dtype, use_default=True)
                        )
                    else:
                        raise RuntimeError(
                            f"Input shape is dynamic but shapes are not provided as dictionary (found: {self.shape})"
                        )

                else:
                    raise ValueError(
                        "Requested an example tensor from a dynamic shaped input but did not specific which profile field to use."
                    )
            raise

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.