Source code for polygraphy.common.struct

# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

from polygraphy import mod, util
from polygraphy.common.interface import TypedDict
from polygraphy.datatype import DataType
from polygraphy.json import Decoder, Encoder, add_json_methods

class BoundedShape(list):
    Represents a shape with min/max bounds.

    def __init__(self, shape, min=None, max=None):
        self.min = min
        self.max = max

    def __repr__(self):
        return f"BoundedShape({list(self)}, min={self.min}, max={self.max})"

class MetadataTuple:
    def __init__(self, dtype, shape, docstring):
        self.dtype = dtype
        self.shape = shape
        self.docstring = docstring

    def dtype(self):
        return self._dtype

    def dtype(self, new):
        self._dtype = DataType.from_dtype(new) if new is not None else None

    def __iter__(self):
        yield from [self.dtype, self.shape]

    def __repr__(self):
        return f"MetadataTuple({self.dtype}, {self.shape}, {self.docstring})"

    def __str__(self):
        ret = ""
        meta_items = []
        if self.dtype is not None:
        if self.shape is not None:
        if self.docstring is not None:
        if meta_items:
            ret += "[" + ", ".join(meta_items) + "]"
        return ret

    def __eq__(self, other):
        return self.shape == other.shape and self.dtype == other.dtype

[docs] @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 """
[docs] @staticmethod def from_feed_dict(feed_dict): """ Constructs a new TensorMetadata using information from the provided feed_dict. Args: feed_dict (OrderedDict[str, Union[numpy.ndarray, torch.tensor]]): A mapping of input tensor names to corresponding input arrays. Returns: TensorMetadata """ meta = TensorMetadata() for name, arr in feed_dict.items(): meta.add(name, util.array.dtype(arr), util.array.shape(arr)) return meta
[docs] def add(self, name, dtype, shape, min_shape=None, max_shape=None, docstring=None): """ Convenience function for adding entries. Args: name (str): The name of the input. dtype (Any): The data type of the input. This can be any type that can be converted to a Polygraphy DataType. 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. docstring (str): Any additional information associated with a tensor. Returns: The newly added entry. """ self[name] = MetadataTuple( dtype, ( BoundedShape(shape, min=min_shape, max=max_shape) if shape is not None else None ), docstring, ) return self
def __repr__(self): ret = "TensorMetadata()" for name, metadata_tuple in self.items(): (dtype, shape) = metadata_tuple ret += util.make_repr( ".add", name, dtype, list(shape), min_shape=shape.min, max_shape=shape.max, docstring=metadata_tuple.docstring, )[0] return ret def __str__(self): sep = ",\n " elems = [f"{name} {meta_tuple}".strip() for name, meta_tuple in self.items()] return "{" + sep.join(elems) + "}"
[docs] @mod.export() @add_json_methods("formatted array") class FormattedArray: """ [EXPERIMENTAL, UNTESTED] This API is experimental and untested and may be significantly modified in future releases. Use with caution! Representes an array whose semantic shape differs from its physical size in memory. For example, consider an ``NCHW`` tensor of shape ``(1, 3, 28, 28)``. If we use a vectorized format like ``N(C/4)HW4``, then the physical size of the array would be ``(1, 1, 28, 28 * 4)`` since the channel dimension would be padded to a multiple of 4. However, we still need a way to keep track of the semantic shape for things like shape inference. This class provides a mechanism to specify the shape of an array independently of the underlying array. """ def __init__(self, array, shape): """ Args: array (Union[np.ndarray, polygraphy.cuda.DeviceView]): The array. In most cases, this will be a raw byte-array. shape (Sequence[int]): The semantic shape of the data. """ self.array = array self.shape = shape
@Encoder.register(FormattedArray) def encode(farray): return { "array": farray.array, "shape": farray.shape, } @Decoder.register(FormattedArray) def decode(dct): return FormattedArray(dct["array"], dct["shape"])