Data Structures

Module: polygraphy.common

class TensorMetadata(dct=None)[source]

Bases: Interface

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
static from_feed_dict(feed_dict)[source]

Constructs a new TensorMetadata using information from the provided feed_dict.

Parameters:

feed_dict (OrderedDict[str, Union[numpy.ndarray, torch.tensor]]) – A mapping of input tensor names to corresponding input arrays.

Returns:

TensorMetadata

add(name, dtype, shape, min_shape=None, max_shape=None, docstring=None)[source]

Convenience function for adding entries.

Parameters:
  • 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.

class FormattedArray(array, shape)[source]

Bases: object

[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.

Parameters:
  • 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.

static from_json(src)

Decode a JSON object and create an instance of this class.

Parameters:

src (str) – The JSON representation of the object

Returns:

The decoded instance

Return type:

FormattedArray

Raises:

PolygraphyException – If the JSON cannot be decoded to an instance of FormattedArray

static load(src)

Loads an instance of this class from a JSON file.

Parameters:

src (Union[str, file-like]) – The path or file-like object to read from.

Returns:

The decoded instance

Return type:

FormattedArray

Raises:

PolygraphyException – If the JSON cannot be decoded to an instance of FormattedArray

save(dest)

Encode this instance as a JSON object and save it to the specified path or file-like object.

Parameters:

dest (Union[str, file-like]) – The path or file-like object to write to.

to_json()

Encode this instance as a JSON object.

Returns:

A JSON representation of this instance.

Return type:

str