Class _DLTensorShapeProxy

python/tvm/relax/expr.py:393–460

A proxy object for unpacking the shape from DLTensor Exposes accessors for the `DLTensor::shape` field. Accessing these fields will produce `relax.Call` expressions, representing the field's runtime value. If the datatype of the tensor is known at compile-time, the `relax.Call` wi

class _DLTensorShapeProxy(tvm.runtime.ObjectConvertible):
    """A proxy object for unpacking the shape from DLTensor

    Exposes accessors for the `DLTensor::shape` field.  Accessing
    these fields will produce `relax.Call` expressions, representing
    the field's runtime value.  If the datatype of the tensor is known
    at compile-time, the `relax.Call` will be normalized into a
    `relax.PrimValue`, with no runtime cost.

    Parameters
    ----------
    tensor: relax.Expr

        The relax tensor (or a variable referring to a relax tensor),
        whose runtime shape is being inspected.
    """

    def __init__(self, tensor):
        self.tensor = tensor

    def asobject(self):
        """Provide expected in error message

        This method is called when `_DLTensorShapeProxy` is used in a
        context that requires a `relax.Expr`.  This usage is not
        supported, and raising an error here can provide suggested
        fixes that are not present in the default error message from
        `tvm.runtime.convert`.
        """
        raise TypeError(
            f"{self.tensor}.shape cannot be converted to a relax expression, "
            f"and should be used as a proxy object to access the runtime shape of the DLTensor. "
            f"The DLTensor::ndim field can be accessed as len({self.tensor}), "
            f"and the DLTensor::shape array can be accessed as {self.tensor}.shape[i]"
        )

    def __getitem__(self, axis: int | PrimExpr | Expr) -> Expr:
        """Returns the extent of a tensor axis

        Parameters
        ----------
        axis: int | PrimExpr | Expr

            The tensor axis whose extent should be returned.  For ease
            of use, any python integers or TIR expressions are
            converted to `relax.Expr`.

        Returns
        -------
        extent: Expr

            The extent of the tensor's axis.
        """

        if not isinstance(axis, tvm.relax.Expr):
            axis = tvm.relax.PrimValue(axis)

        if axis.struct_info_ is not None and not isinstance(