Class TensorMetadata

tools/Polygraphy/polygraphy/common/struct.py:66–131

An OrderedDict[str, MetadataTuple] that maps input names to their data types and shapes. Shapes may include negative values, ``None``, or strings to indicate dynamic dimensions. Example: :: shape = tensor_meta["input0"].shape dtype = tensor_meta["input0"].dtype

@mod.export()
class TensorMetadata(TypedDict(lambda: str, lambda: MetadataTuple)):
    """
    An OrderedDict[str, MetadataTuple] that maps input names to their data types and shapes.

    Shapes may include negative values, ``None``, or strings to indicate dynamic dimensions.

    Example:
    ::

        shape = tensor_meta["input0"].shape
        dtype = tensor_meta["input0"].dtype
    """

    @staticmethod
    def from_feed_dict(feed_dict):
        """
        Constructs a new TensorMetadata using information from the provided feed_dict.

        Args:
            feed_dict (OrderedDict[str, numpy.ndarray]):
                    A mapping of input tensor names to corresponding input NumPy arrays.

        Returns:
            TensorMetadata
        """
        meta = TensorMetadata()
        for name, arr in feed_dict.items():
            meta.add(name, arr.dtype, arr.shape)
        return meta

    def add(self, name, dtype, shape, min_shape=None, max_shape=None):
        """
        Convenience function for adding entries.

        Args:
            name (str): The name of the input.
            dtype (numpy.dtype): The data type of the input.
            shape (Sequence[Union[int, str]]]):
                    The shape of the input. Dynamic dimensions may
                    be indicated by negative values, ``None``, or a string.

            min_shape (Sequence[int]):
                    The minimum valid shape for the input.
                    If provided, this shape should not include any dynamic dimensions.
            max_shape (Sequence[int]):
                    The maximum valid shape for the input.
                    If provided, this shape should not include any dynamic dimensions.

        Returns:
            The newly added entry.
        """
        self[name] = MetadataTuple(
            dtype, BoundedShape(shape, min=min_shape, max=max_shape) if shape is not None else None
        )
        return self

    def __repr__(self):
        ret = "TensorMetadata()"