tensorrt-cu12-bindings 10.13.3.9.post1__cp312-none-win_amd64.whl

tensorrt_bindings/plugin/_tensor.py

@@ -0,0 +1,1128 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import tensorrt as trt
+from typing import Tuple, Union
+import numpy as np
+from ._export import public_api, IS_AOT_ENABLED
+from abc import ABC, abstractmethod
+
+class SymExpr(ABC):
+    
+    @abstractmethod
+    def _op(self, op, other):
+        pass
+
+    @abstractmethod
+    def __add__(self, other):
+        pass
+
+    @abstractmethod
+    def __sub__(self, other):
+        pass
+
+    @abstractmethod
+    def __mul__(self, other):
+        pass
+
+    @abstractmethod
+    def __floordiv__(self, other):
+        pass
+
+    @abstractmethod
+    def __eq__(self, other):
+        pass
+
+    @abstractmethod
+    def __lt__(self, other):
+        pass
+
+    @abstractmethod
+    def __repr__(self):
+        pass
+
+    @property
+    @abstractmethod
+    def is_constant(self) -> bool:
+        pass
+
+    @property
+    @abstractmethod
+    def constant_value(self) -> int:
+        pass
+    
+    # Evaluate the underlying trt.IDimensionExpr, if so done lazily
+    @property
+    @abstractmethod
+    def _expr(self):
+        pass
+
+class SymIntExprMeta(type(SymExpr)):
+    pass
+
+class SymIntExpr(SymExpr, metaclass=SymIntExprMeta):
+    """
+    Symbolic integer (scalar) expression
+    """
+    _exprBuilder = None  # trt.IExprBuilder instance. Populated when a shape-calculation context is entered.
+
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "SymIntExpr"] = None):
+        """
+        Args:
+            value (Union[int, trt.IDimensionExpr, SymIntExpr], optional): Constant or another symbolic expression. Defaults to creating a fake shape expression.
+        """
+        self._int_expr = None
+        if isinstance(value, int):
+            if SymIntExpr._exprBuilder is None:
+                self._int_expr = None
+            else:
+                self._int_expr = SymIntExpr._exprBuilder.constant(value)
+        elif isinstance(value, trt.IDimensionExpr):
+            self._int_expr = value
+        elif isinstance(value, SymIntExpr):
+            self._int_expr = value._int_expr
+
+    def _op(self, op: trt.DimensionOperation, other: Union[int, "SymIntExpr"]):
+        if isinstance(other, int):
+            other = SymIntExpr(other)
+        return SymIntExpr(SymIntExpr._exprBuilder.operation(op, self._expr, other._expr))
+
+    # Binary operations for +, -, *, //, ==. <
+    # Those for ceil_div, max and min are provided as top-level functions of tensorrt.plugin
+    def __add__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.SUM, other)
+
+    def __sub__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.SUB, other)
+
+    def __mul__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.PROD, other)
+
+    def __floordiv__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.FLOOR_DIV, other)
+
+    def __eq__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.EQUAL, other)
+
+    def __lt__(self, other: Union[int, "SymIntExpr"]):
+        return self._op(trt.DimensionOperation.LESS, other)
+
+    def __repr__(self):
+        if self._is_dummy:
+            return f"FakeSymIntExpr[id={id(self)}]"
+        elif not self.is_constant:
+            return f"SymIntExpr[id={id(self)}]"
+        return f"SymIntExpr[{self._expr.get_constant_value()}]"
+
+    @property
+    def is_constant(self) -> bool:
+        """
+        `True` if this integer expression is a build-time constant, `False` otherwise.
+
+        Raises:
+            RuntimeError: For fake :class:`SymIntExpr`\s. Check :attr:`is_fake` to determine accessibility.
+        """
+        if self._is_dummy:
+            raise RuntimeError(
+                "Not accessible for fake 'SymIntExpr's. Check is_fake to determine accessibility."
+            )
+        return self._expr.is_constant()
+
+    def constant_value(self) -> int:
+        """
+        Return value of the constant integer expression.
+
+        Raises:
+            RuntimeError: For non-constant integer expressions. Check :attr:`is_constant` to determine accessibility.
+        """
+        if not self.is_constant:
+            raise RuntimeError(
+                "Not accessible for non-constant integer expressions. Check is_constant to determine accessibility."
+            )
+        return self._expr.get_constant_value()
+    
+    # Evaluate the underlying trt.IDimensionExpr, if so done lazily
+    @property
+    def _expr(self):
+        return self._int_expr
+
+    def _clone(self):
+        return SymIntExpr(self + 0)
+
+class SymExprImpl(trt.ISymExpr):
+    def __init__(self, expr: SymIntExpr):
+        trt.ISymExpr.__init__(self)
+        self.type = trt.PluginArgType.INT
+        if isinstance(expr, SymInt32):
+            self.dtype = trt.PluginArgDataType.INT32
+        elif isinstance(expr, SymInt16):
+            self.dtype = trt.PluginArgDataType.INT16
+        elif isinstance(expr, SymInt8):
+            self.dtype = trt.PluginArgDataType.INT8
+        else:
+            raise ValueError(f"Unknown SymIntExpr type {type(expr)}")
+        
+        self.expr = expr._expr
+@public_api()
+class SymInt32(SymIntExpr):
+    """
+    Symbolic expression for a 32-bit integer
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, SymIntExpr] = None):
+        super().__init__(value)
+
+    def __call__(self):
+        return SymExprImpl(self)
+@public_api()
+class SymInt8(SymIntExpr):
+    """
+    Symbolic expression for an 8-bit integer
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "SymIntExpr"] = None):
+        super().__init__(value)
+
+@public_api()
+class SymInt16(SymIntExpr):
+    """
+    Symbolic expression for a 16-bit integer
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "SymIntExpr"] = None):
+        super().__init__(value)
+
+# Symbolic expression for a given dimension of a tensor
+@public_api()
+class ShapeExpr(SymInt32):
+    """
+    Symbolic expression for single dimension of a tensor
+    """
+    def __init__(self, value: Union[int, trt.IDimensionExpr, "ShapeExpr", SymIntExpr] = None):
+        """
+        Args:
+            value (Union[int, trt.IDimensionExpr, ShapeExpr, SymIntExpr], optional): Constant or another symbolic expression. Defaults to creating a fake shape expression.
+        """
+        super().__init__(value)
+        self._exprBuilder = SymIntExpr._exprBuilder
+        self._is_dummy = False
+        self._is_size_tensor = False
+        if value is None:
+            self._is_dummy = True
+        elif isinstance(value, int):
+            if self._exprBuilder is None:
+                self._is_dummy = True
+        elif isinstance(value, ShapeExpr):
+            self._is_dummy = value._is_dummy
+            self._is_size_tensor = value._is_size_tensor
+        elif isinstance(value, SymIntExpr):
+            pass
+
+    def _op(self, op: trt.DimensionOperation, other: Union[int, "ShapeExpr"]):
+        if self._is_size_tensor:
+            raise ValueError("It is not permitted to perform binary operations on size tensor expressions") # trt limitation
+        if self._is_dummy:
+            return ShapeExpr()
+        return ShapeExpr(super()._op(op, other))
+
+    def __repr__(self):
+        if self._is_dummy:
+            return f"FakeShapeExpr[id={id(self)}]"
+        elif not self.is_constant:
+            return f"ShapeExpr[id={id(self)}]"
+        return f"ShapeExpr[{self._expr.get_constant_value()}]"
+
+    # A ShapeExpr may be "fake" when it is accessed in a non-shape calculation context. Fake `ShapeExpr`s are externally indistinguishable unless `is_constant` or `constant_value` is required.
+    # Therefore, constant checks/access must occur conditionally after evaluating `is_fake`.
+    @property
+    def is_fake(self) -> bool:
+        """
+        A ShapeExpr may be "fake" when it is accessed in a non-shape calculation context.
+        Fake `ShapeExpr`s are externally indistinguishable unless `is_constant` or `constant_value` is required.
+        """
+        return self._is_dummy
+
+    @property
+    def is_size_tensor(self) -> bool:
+        """
+        `True` if this represents a size tensor, `False` otherwise.
+        """
+        return self._is_size_tensor
+
+    @property
+    def is_constant(self) -> bool:
+        """
+        `True` if this shape expression is a build-time constant, `False` otherwise.
+
+        Raises:
+            RuntimeError: For fake :class:`ShapeExpr`\s. Check :attr:`is_fake` to determine accessibility.
+        """
+        if self._is_dummy:
+            raise RuntimeError(
+                "Not accessible for fake 'ShapeExpr's. Check is_fake to determine accessibility."
+            )
+        return super().is_constant
+
+    def constant_value(self) -> int:
+        """
+        Return value of the constant shape expression.
+
+        Raises:
+            RuntimeError: For non-constant shape expressions. Check :attr:`is_constant` to determine accessibility.
+        """
+        if not self.is_constant:
+            raise RuntimeError(
+                "Not accessible for non-constant shape expressions. Check is_constant to determine accessibility."
+            )
+        return super().constant_value()
+
+    def _clone(self):
+        ret = ShapeExpr(self + 0)
+        ret._is_dummy = self._is_dummy
+        ret._is_size_tensor = self._is_size_tensor
+        return ret
+
+@public_api()
+class SizeTensorShapeExpr(ShapeExpr):
+    """
+    Extends :class:`ShapeExpr`
+
+    A shape expression that represent a size tensor
+
+    """
+    def __init__(self, size_tensor_desc: "SizeTensorDesc"):
+        """
+        .. note:: It is recommended to use :attr:`SizeTensorDesc.expr` to get a :class:`SizeTensorShapeExpr` representing a size tensor
+        """
+        super().__init__()
+        self._is_size_tensor = True
+        self._is_dummy = size_tensor_desc.opt.is_fake
+        self._size_tensor_desc = size_tensor_desc
+
+    def _op(self, op: trt.DimensionOperation, other: Union[int, "ShapeExpr"]):
+        raise ValueError("It is not permitted to perform binary operations on size tensor expressions") # TRT limitation
+
+    @property
+    def is_constant(self):
+        if self._is_dummy:
+            raise RuntimeError(
+                "Not accessible for fake 'ShapeExpr's. Check is_fake to determine accessibility."
+            )
+        return False
+
+    @property
+    def _expr(self):
+        if self._int_expr is not None:
+            return self._int_expr
+
+        self._int_expr = super()._exprBuilder.declare_size_tensor(self._size_tensor_desc.index, self._size_tensor_desc.opt._expr, self._size_tensor_desc.upper_bound._expr)
+        return self._int_expr
+    
+    def __repr__(self):
+        return f"ShapeExpr[is_size_tensor = True, id={id(self)}]"
+
+def _from_scalar(s):
+    if isinstance(s, int):
+        return SymInt32(s)
+    elif isinstance(s, float):
+        raise ValueError("Float symbolic expressions are not supported")
+    else:
+        raise ValueError(f"Unsupported type: '{type(s)}'")
+
+# Iterable holding `ShapeExpr`s
+@public_api()
+class SymExprs:
+    def __init__(self, length: int):
+        """
+        Iterable holding symbolic expressions
+
+        Args:
+            length (int): Number of dimensions of the tensor
+        """
+        self._length = length
+        self._exprs = [None] * length
+
+    @classmethod
+    def from_tuple(cls, shape_exprs: Tuple[Union[SymExpr, int]]) -> "SymExprs":
+        """
+        Args:
+            shape_exprs (Tuple[Union[SymExpr, int]]): Tuple to construct :class:`SymExprs` from
+        """
+
+        shape_exprs_ = tuple([e if isinstance(e, SymExpr) else _from_scalar(e) for e in shape_exprs])
+        inst = cls(len(shape_exprs_))
+        inst._exprs = list(shape_exprs_)
+        return inst
+
+    def __iter__(self):
+        return iter(self._exprs)
+
+    def __getitem__(self, index):
+        return self._exprs[index]
+
+    def __len__(self):
+        return self._length
+
+    def __setitem__(self, index, expr):
+        if index >= self._length:
+            raise IndexError("Index out of range")
+        
+        if not isinstance(expr, SymExpr):
+            expr = _from_scalar(expr)
+
+        self._exprs[index] = expr
+
+    def __repr__(self):
+        return f"SymExprs[{', '.join([s.__repr__() for s in self._exprs])}]"
+
+@public_api()
+class ShapeExprs(SymExprs):
+    def __init__(self, length, _is_dummy = False):
+        """
+        Iterable holding :class:`ShapeExpr`\s, representing a tensor shape
+
+        Args:
+            length (int): Number of dimensions of the tensor
+        """
+        if length > trt.Dims.MAX_DIMS:
+            raise ValueError(f"ShapeExprs can only support up to trt.Dims.MAX_DIMS = {trt.Dims.MAX_DIMS} dimensions. {length} given.")
+        
+        super().__init__(length)
+
+        self._is_dummy = _is_dummy
+        if _is_dummy:
+            self._exprs = [ShapeExpr()] * length
+
+    @classmethod
+    def from_tuple(cls, shape_exprs: Tuple[Union[ShapeExpr, int]]) -> "ShapeExpr":
+        """
+        Args:
+            shape_exprs (Tuple[Union[ShapeExpr, int]]): Tuple to construct :class:`ShapeExprs` from
+        """
+
+        shape_exprs_ = tuple([e if isinstance(e, ShapeExpr) else ShapeExpr(e) for e in shape_exprs])
+        inst = cls(len(shape_exprs_))
+        inst._exprs = list(shape_exprs_)
+        return inst
+    
+    def numel(self) -> ShapeExpr:
+        """
+        Returns a symbolic expression for the number of elements
+        """
+        ret = ShapeExpr(1)
+        for s in self._exprs:
+            ret *= s
+        return ret
+
+    def __setitem__(self, index, value):
+        if index >= self._length:
+            raise IndexError("Index out of range")
+        
+        if not isinstance(value, ShapeExpr):
+            if not isinstance(value, int):
+                raise ValueError(f"Value should be int or ShapeExpr. Got '{type(value)}'")
+            value = ShapeExpr(value)
+
+        self._exprs[index] = value
+    
+    def __repr__(self):
+        return f"ShapeExprs[{', '.join([s.__repr__() for s in self._exprs])}]"
+
+    def _clone(self):
+        ret = ShapeExprs.from_tuple((e._clone() for e in self._exprs))
+        ret._is_dummy = self._is_dummy
+        return ret
+    
+@public_api()
+class SymIntExprs(SymExprs):
+    def __init__(self, length):
+        """
+        Iterable holding :class:`SymIntExpr`\s
+
+        Args:
+            length (int): Number of symbolic expressions in the iterable
+        """   
+        super().__init__(length)
+
+    @classmethod
+    def from_tuple(cls, shape_exprs: Tuple[Union[SymIntExpr, int]]) -> "SymIntExpr":
+        """
+        Args:
+            shape_exprs (Tuple[Union[SymIntExpr, int]]): Tuple to construct :class:`SymIntExprs` from
+        """
+
+        shape_exprs_ = tuple([e if isinstance(e, SymIntExpr) else SymIntExpr(e) for e in shape_exprs])
+        inst = cls(len(shape_exprs_))
+        inst._exprs = list(shape_exprs_)
+        return inst
+
+    def __setitem__(self, index, value):
+        if index >= self._length:
+            raise IndexError("Index out of range")
+        
+        if not isinstance(value, SymIntExpr):
+            if not isinstance(value, int):
+                raise ValueError(f"Value should be int or SymIntExpr. Got '{type(value)}'")
+            value = SymIntExpr(value)
+
+        self._exprs[index] = value
+    
+    def __repr__(self):
+        return f"SymIntExprs[{', '.join([s.__repr__() for s in self._exprs])}]"
+
+# Numerical representation of a tensor shape
+@public_api()
+class Shape:
+    """
+    Numerical representation of a tensor shape
+    """
+    def __init__(
+        self, tensor_desc: Union[Tuple[int], trt.DynamicPluginTensorDesc, trt.PluginTensorDesc] = None
+    ):
+        self._is_dynamic = None  # set lazily
+        if isinstance(tensor_desc, trt.DynamicPluginTensorDesc):
+            self._length = len(tensor_desc.desc.dims)
+            self._shapes = tensor_desc.desc.dims
+            self._desc = tensor_desc
+        elif isinstance(tensor_desc, trt.PluginTensorDesc):
+            self._length = len(tensor_desc.dims)
+            self._shapes = tensor_desc.dims
+        elif isinstance(tensor_desc, tuple):
+            self._shapes = trt.Dims(tensor_desc)
+            self._length = len(self._shapes)
+        elif tensor_desc is None:
+            self._length = 0
+            self._shapes = trt.Dims(0)
+        else:
+            raise ValueError("Unsupported type used for constructing trt.plugin.Shape! tensor_desc must be a Tuple[int], trt.DynamicPluginTensorDesc, or trt.PluginTensorDesc")
+
+    def numel(self) -> int:
+        """
+        Number of elements contained
+
+        Raises:
+            ValueError: When :attr:`is_dynamic` is `True`
+        """
+        if self.is_dynamic:
+            raise ValueError("Shape has at least one dynamic dimension.")
+        return int(np.prod(self._shapes))
+
+    def __iter__(self):
+        yield from self._shapes
+
+    def __getitem__(self, index):
+        return self._shapes[index]
+
+    def __len__(self):
+        return self._length
+
+    def __str__(self):
+        return "Shape" + str(tuple(self))
+
+    @property
+    def is_dynamic(self) -> bool:
+        """
+        `True` if this tensor has at least one dynamic dimension, `False` otherwise.
+        """
+        if self._is_dynamic is not None:
+            return self._is_dynamic
+
+        self._is_dynamic = False
+        for d in self._shapes:
+            if d == -1:
+                self._is_dynamic = True
+
+        return self._is_dynamic
+
+    @property
+    def opt(self) -> Tuple[int]:
+        """
+        Optimum value of dimensions specified for auto-tuning.
+        """
+        if not self.is_dynamic:
+            raise ValueError("opt property is only accessible if is_dynamic is true")
+        if not hasattr(self, "_desc"):
+            raise AttributeError(
+            "Shape object has at least one dynamic dimension, but no information is available on 'opt' property."
+            )
+        return tuple(self._desc.opt)
+
+    @property
+    def min(self) -> Tuple[int]:
+        """
+        Lower bounds on tensor's dimensions.
+        """
+        if not self.is_dynamic:
+            raise ValueError("min property is only accessible if is_dynamic is true")
+        if not hasattr(self, "_desc"):
+            raise AttributeError(
+            "Shape object has at least one dynamic dimension, but no information is available on 'min' property."
+            )
+        return tuple(self._desc.min)
+
+    @property
+    def max(self) -> Tuple[int]:
+        """
+        Upper bounds on tensor's dimensions.
+        """
+        if not self.is_dynamic:
+            raise ValueError("max property is only accessible if is_dynamic is true")
+        if not hasattr(self, "_desc"):
+            raise AttributeError(
+            "Shape object has at least one dynamic dimension, but no information is available on 'max' property."
+            )
+        return tuple(self._desc.max)
+
+    def __setitem__(self, index, val):
+        if index >= self._length:
+            raise IndexError("Index out of range")
+        self._shapes[index] = val
+
+    def _clone(self):
+        ret = Shape()
+        ret.__dict__.update(self.__dict__)
+        return ret
+
+
+# Descriptor for a tensor
+# A `TensorDesc` never contains nor refers to any tensor data.
+@public_api()
+class TensorDesc:
+    """
+    Descriptor for a tensor
+    A `TensorDesc` never contains nor refers to any tensor data.
+    """
+    def __init__(self, shape_expr: ShapeExprs = None, dtype: trt.DataType = None, format: trt.TensorFormat = None, scale: float = None):
+        """
+        Args:
+            shape_expr (ShapeExprs): The data with which to initialize the tensor.
+            dtype (trt.DataType): The data type of the tensor.
+            format (trt.TensorFormat): Format (layout) of the tensor.
+            scale (float): Scale for INT8 data type.
+
+        .. code-block:: python
+            :linenos:
+            :caption: Creates a TensorDesc with constant shape expressions
+
+            tensor = trt.TensorDesc((10, 2, 32, 32), dtype=trt.float32)
+
+        .. code-block:: python
+            :linenos:
+            :caption: Creates a TensorDesc from shape expression of another TensorDesc
+
+            tensor = trt.from_shape_expr(other.shape_expr, dtype=trt.float32)
+        """
+
+        # `TensorDesc` may or may not have `Shape` information but always has symbolic shape expressions and dtype
+        self._shape_expr = shape_expr
+        self._dtype = dtype
+
+        # `shape`, `format`, and `scale` are only accessible if `has_shape`. Presently, this would be inside autotune.
+        self._shape = None
+        self._format = format
+        self._scale = scale
+
+        self._aliased_to = None
+        self._immutable = False
+
+    def numel(self) -> int:
+        """
+        Returns:
+            Returns an int with the number of elements of the tensor.
+
+        .. warning::
+            Should only be called when TensorDesc.has_shape is true. If a symbolic expression for the number of elements is required, query TensorDesc.shape_expr.numel().
+        """
+        if not self.has_shape:
+            raise ValueError(
+                "TensorDesc has no shape information available at this stage. Inspect TensorDesc.has_shape to determine availability."
+            )
+        return int(np.prod(self.shape))
+
+    @property
+    def ndim(self) -> int:
+        """
+        Number of dimensions
+        """
+        return len(self._shape_expr)
+
+    @property
+    def is_size_tensor(self):
+        return False
+
+    # Return a `TensorDesc` that has identical properties to `self` but is mutable
+    def like(self) -> "TensorDesc":
+        """
+        Returns:
+            Returns a TensorDesc which has identical properties to this tensor, and is mutable.
+
+        .. code-block:: python
+            :linenos:
+            :caption: Communicate that output tensor has identical properties to the input tensor
+
+            @tensorrt.plugin.register("my::plugin")
+            def _(inp: tensorrt.plugin.TensorDesc) -> tensorrt.plugin.TensorDesc:
+                return inp.like()
+        """
+        cloned = self._clone()
+        cloned._immutable = False
+        return cloned
+
+    # Return a `TensorDesc` that has identical properties to `self` AND is aliased to `self` (would result in a `Tensor` during enqueue sharing the same data buffer)
+    def aliased(self) -> "TensorDesc":
+        """
+        Returns:
+            Returns a TensorDesc which has identical properties and is aliased to this tensor (would result in a `Tensor` during enqueue sharing the same data buffer).
+            Returned TensorDesc is immutable.
+
+        .. code-block:: python
+            :linenos:
+            :caption: Communicate that output tensor has identical properties to the input tensor
+
+            @tensorrt.plugin.register("my::plugin")
+            def _(inp: tensorrt.plugin.TensorDesc) -> tensorrt.plugin.TensorDesc:
+                return inp.aliased()
+        """
+        cloned = self._clone()
+        cloned._immutable = False
+        cloned._aliased_to = self
+        cloned._immutable = True
+        return cloned
+
+    def _clone(self) -> "TensorDesc":
+        cloned = TensorDesc()
+        cloned.__dict__.update(self.__dict__)
+        cloned._immutable = False
+        cloned._shape_expr = self._shape_expr._clone()
+        if self._shape is not None:
+            cloned._shape = self._shape._clone()
+        cloned._immutable = True
+        return cloned
+
+    def get_aliased(self) -> "TensorDesc":
+        """
+        Returns:
+            Returns a TensorDesc for the tensor which this tensor is aliased to. Returns None is this tensor is not aliased to any other tensor.
+        """
+        return self._aliased_to
+
+    def _validate_has_shape(self) -> None:
+        if not self.has_shape:
+            raise ValueError(
+                "TensorDesc has no shape information available at this stage. Inspect TensorDesc.has_shape to determine availability."
+            )
+
+    def _validate_not_immutable(self):
+        if hasattr(self, "_immutable") and self._immutable:
+            raise ValueError("Cannot modify immutable TensorDesc")
+
+    @property
+    def shape_expr(self) -> ShapeExprs:
+        """
+        Symbolic expressions for the tensor shape.
+        """
+        return self._shape_expr
+
+    @property
+    def dtype(self) -> trt.DataType:
+        """
+        Data type of the tensor.
+        """
+        return self._dtype
+
+    @property
+    def shape(self) -> Shape:
+        """
+        The (concrete) shape of the tensor.
+
+        .. warning::
+            Only accessible when TensorDesc.has_shape is true.
+        """
+        self._validate_has_shape()
+        return self._shape
+
+    @property
+    def format(self) -> trt.TensorFormat:
+        """
+        The format of the tensor.
+
+        .. warning::
+            Only accessible when TensorDesc.has_shape is true.
+        """
+        self._validate_has_shape()
+        return self._format
+
+    @property
+    def scale(self) -> float:
+        """
+        Scale for INT8 data type.
+
+        .. warning::
+            Only accessible when TensorDesc.has_shape is true.
+        """
+        self._validate_has_shape()
+        return self._scale
+
+
+    @shape_expr.setter
+    def shape_expr(self, value):
+        self._shape_expr = value
+
+    @dtype.setter
+    def dtype(self, value):
+        self._dtype = value
+
+    @shape.setter
+    def shape(self, value):
+        self._validate_not_immutable()
+        self._shape = value
+
+    @format.setter
+    def format(self, value):
+        self._validate_not_immutable()
+        self._format = value
+
+    @scale.setter
+    def scale(self, value):
+        self._validate_not_immutable()
+        self._scale = value
+
+    @property
+    def is_aliased(self) -> bool:
+        """
+        True if this tensor is aliased to another tensor, False otherwise.
+        """
+        return self._aliased_to is not None
+
+    @property
+    def has_shape(self) -> bool:
+        """
+        True if this tensor has concrete shape information, False otherwise.
+        """
+        return self._shape is not None
+
+    @property
+    def is_dynamic(self) -> bool:
+        """
+        `True` if this tensor has at least one dynamic dimension, `False` otherwise.
+        """
+        if not self.has_shape:
+            raise ValueError(
+                "TensorDesc has no shape information available at this stage. Inspect TensorDesc.has_shape to determine availability."
+            )
+        return self.shape.is_dynamic
+
+    @property
+    def has_shape_expr(self) -> bool:
+        """
+        True if this tensor has symbolic shape expressions, False otherwise.
+        """
+        return self.shape_expr is not None
+
+    def __setattr__(self, name, value):
+        if hasattr(self, "_immutable") and self._immutable and name != "_immutable":
+            raise ValueError("Cannot modify immutable TensorDesc properties")
+        super().__setattr__(name, value)
+
+@public_api()
+class SizeTensorDesc(TensorDesc):
+    """
+    Extends :class:`TensorDesc`
+
+    Descriptor for a size tensor: a scalar of either INT32 or INT64 data type used to express the extent of a data-dependent dimension.
+    """
+    def __init__(self, opt: ShapeExpr, upper_bound: ShapeExpr):
+        """
+        Args:
+            opt (ShapeExpr): Symbolic expression for the extent of this size tensor to use in the autotune process of the engine build
+            upper_bound (ShapeExpr): Symbolic expression for the upper-bound of this size tensor
+
+        .. note:: It is recommended to construct a size tensor using :func:`size_tensor` instead of using this constructor directly
+        """
+        super().__init__(ShapeExprs(0), trt.int32)
+        self._opt = opt
+        self._upper_bound = upper_bound
+        self._index = None
+        self._expr = SizeTensorShapeExpr(self)
+
+    @property
+    def is_size_tensor(self):
+        return True
+
+    @property
+    def opt(self) -> ShapeExpr:
+        """
+        Symbolic expression for the extent of this size tensor to use in the autotune process of the engine build
+        """
+        return self._opt
+
+    @property
+    def upper_bound(self) -> ShapeExpr:
+        """
+        Symbolic expression for the upper-bound of this size tensor
+        """
+        return self._upper_bound
+
+    @property
+    def index(self) -> int:
+        """
+        Output index at which this size tensor resides
+        """
+        return self._index
+
+    def _set_index(self, idx: int):
+        self._index = idx
+
+    def expr(self) -> SizeTensorShapeExpr:
+        """
+        Symbolic expression for this size tensor
+        """
+        return self._expr
+
+
+# A tensor representation that carries data
+@public_api()
+class Tensor:
+    """
+    Representation of a tensor that carries data
+
+    :class:`Tensor` objects are strictly *descriptors* of a tensor with an underlying data buffer. `tensorrt.plugin` does not provide any APIs that perform standard data-altering operations on :class:`Tensor`\s.
+
+    Supports `__cuda_array_interface__` for interoperability with other frameworks.
+
+    """
+    def __init__(self):
+        self._data_ptr = None
+        self._shape = None
+        self._format = None
+        self._dtype = None
+        self._scale = None
+        self._strides = None
+
+        self._aliased_to = None
+        self._stream = None
+        self._read_only = None
+        self._immutable = False
+
+    @property
+    def ndim(self) -> int:
+        """
+        Number of dimensions
+        """
+        return len(self._shape)
+
+    @property
+    def data_ptr(self) -> int:
+        """
+        Pointer to the data buffer of this tensor
+        """
+        return self._data_ptr
+
+    @property
+    def dtype(self) -> trt.DataType:
+        """
+        Data type of the tensor.
+        """
+        return self._dtype
+
+    @property
+    def shape(self) -> Shape:
+        """
+        The (concrete) shape of the tensor.
+        """
+        return self._shape
+
+    @property
+    def format(self) -> trt.TensorFormat:
+        """
+        The format of the tensor.
+        """
+        return self._format
+
+    @property
+    def scale(self) -> float:
+        """
+        Scale for INT8 data type.
+        """
+        return self._scale
+
+    @property
+    def strides(self) -> Tuple[int]:
+        """
+        Strides of this tensor.
+        """
+        return self._strides
+
+    @data_ptr.setter
+    def data_ptr(self, value):
+        self._data_ptr = value
+
+    @dtype.setter
+    def dtype(self, value):
+        self._dtype = value
+
+    @shape.setter
+    def shape(self, value):
+        self._shape = value
+
+    @format.setter
+    def format(self, value):
+        self._format = value
+
+    @scale.setter
+    def scale(self, value):
+        self._scale = value
+
+    @strides.setter
+    def strides(self, value):
+        self._strides = value
+
+    def numel(self) -> int:
+        """
+        Returns the number of elements of the tensor
+
+        Raises:
+            ValueError: If the tensor has a data-dependent dimension. Examine :attr:`is_data_dependent` to determine whether the tensor is data-dependent.
+
+        Returns:
+            int: Number of elements of the tensor
+        """
+        if self.is_data_dependent:
+            raise ValueError(
+                "Tensor has a data-dependent dimension. Examine Tensor.shape to determine wildcards (representing data-dependent dimensions)."
+            )
+        return int(np.prod(self._shape))
+
+    @property
+    def __cuda_array_interface__(self):
+        if self._dtype in [trt.DataType.BF16, trt.DataType.FP8, trt.DataType.INT4]:
+            raise ValueError(
+                f"Handling {self._dtype} via '__cuda_array_interface__' is not supported"
+            )
+
+        desc = {
+            "shape": tuple(self._shape),
+            "typestr": np.dtype(trt.nptype(self._dtype)).str,
+        }
+        desc["stream"] = self._stream
+        desc["version"] = 3
+        desc["data"] = (
+            self._data_ptr,
+            False,
+        )  # torch does not support read_only flag. Always set to False -- it is user's responsibility to respect implied read-write restriction(s).
+        desc["strides"] = tuple(
+            [s * np.dtype(trt.nptype(self._dtype)).itemsize for s in self._strides]
+        )
+
+        return desc
+
+    def __setattr__(self, name, value):
+        if hasattr(self, "_immutable") and self._immutable and name != "_immutable":
+            raise ValueError("Cannot modify immutable Tensor properties")
+        super().__setattr__(name, value)
+
+    def get_aliased(self) -> "Tensor":
+        """
+        Returns:
+            Returns :class:`Tensor` of the tensor which this tensor is aliased to. Returns None is this tensor is not aliased to any other tensor.
+        """
+        return self._aliased_to
+
+    @property
+    def is_aliased(self):
+        """
+        True if this tensor is aliased to another tensor, False otherwise.
+        """
+        return self._aliased_to is None
+
+    @property
+    def is_data_dependent(self):
+        """
+        True if this tensor contains at least one data-dependent dimension, False otherwise.
+        """
+        return self._shape.is_dynamic
+
+    # Return a `Tensor` which has the same `data_ptr` as `self` but has the provided shape.
+    def aliased(self, shape: Union[Shape, Tuple[int], trt.PluginTensorDesc] = None) -> "Tensor":
+        """
+        Return a :class:`Tensor` which has the same :attr:`data_ptr` as this but has the provided `shape`.
+
+        Args:
+            shape (Union[Shape, Tuple[int], trt.PluginTensorDesc], optional): Required shape of the new tensor (must have the same volume). Defaults to same shape.
+
+        Raises:
+            ValueError: If `shape` is not a supported type or if it does not have the same volume
+        """
+        cloned = Tensor()
+        cloned.__dict__.update(self.__dict__)
+        cloned._immutable = False
+        if isinstance(shape, trt.PluginTensorDesc):
+            cloned._shape = Shape(shape)
+        elif isinstance(shape, Shape):
+            cloned._shape = shape
+        elif isinstance(shape, tuple):
+            desc = trt.PluginTensorDesc()
+            desc.dims = shape
+            desc.type = self._dtype
+            desc.format = self._format
+            desc.scale = self._scale
+            cloned._shape = Shape(desc)
+        elif shape is None:
+            pass
+        else:
+            raise ValueError("Unsupported type for 'shape'")
+
+        # If either the `shape` or self._shape has a wildcard, we allow aliasing
+        if not self.is_data_dependent and cloned.is_data_dependent:
+            if cloned._shape.numel() > self.numel():
+                raise ValueError("Volume of this tensor is less than the provided 'shape'.")
+
+        cloned._aliased_to = self
+        return cloned
+
+if IS_AOT_ENABLED:
+    @public_api()
+    class KernelLaunchParams:
+        """
+        Args:
+            grid_x (Union[int, trt.IDimensionExpr, SymInt32], optional): The grid x dimension. Defaults to 1.
+            grid_y (Union[int, trt.IDimensionExpr, SymInt32], optional): The grid y dimension. Defaults to 1.
+            grid_z (Union[int, trt.IDimensionExpr, SymInt32], optional): The grid z dimension. Defaults to 1.
+            block_x (Union[int, trt.IDimensionExpr, SymInt32], optional): The x dimension of each thread block. Defaults to 1.
+            block_y (Union[int, trt.IDimensionExpr, SymInt32], optional): The y dimension of each thread block. Defaults to 1.
+            block_z (Union[int, trt.IDimensionExpr, SymInt32], optional): The z dimension of each thread block. Defaults to 1.
+            shared_mem (Union[int, trt.IDimensionExpr, SymInt32], optional): Shared-memory per thread block in bytes. Defaults to 0.
+        """
+        def __init__(self,
+                    grid_x: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    grid_y: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    grid_z: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    block_x: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    block_y: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    block_z: Union[int, trt.IDimensionExpr, SymInt32] = 1,
+                    shared_mem: Union[int, trt.IDimensionExpr, SymInt32] = 0):
+            self.grid_x = SymInt32(grid_x)
+            self.grid_y = SymInt32(grid_y)
+            self.grid_z = SymInt32(grid_z)
+            self.block_x = SymInt32(block_x)
+            self.block_y = SymInt32(block_y)
+            self.block_z = SymInt32(block_z)
+            self.shared_mem = SymInt32(shared_mem)
+
+
+        def __setattr__(self, name, value):
+            if name in ["grid_x", "grid_y", "grid_z", "block_x", "block_y", "block_z", "shared_mem"]:
+                self.__dict__[name] = SymInt32(value)
+            else:
+                raise AttributeError(f"KernelLaunchParams object has no attribute '{name}'")