aitune.torch.module.graph_spec | NVIDIA AITune Documentation

Contains GraphSpec which represents a graph specification.

Module Contents

Classes

Name	Description
`GraphSpec`	GraphSpec used to describe a computational graph.

API

class aitune.torch.module.graph_spec.GraphSpec(
    name: str,
    input_spec: aitune.torch.module.sample_metadata.SampleMetadata,
    output_spec: aitune.torch.module.sample_metadata.SampleMetadata,
    forward_signature: aitune.torch.module.forward_signature.ForwardSignature,
    dynamic_shapes: aitune.torch.dynamic_shapes.DynamicShapes = dict()
)

Dataclass

GraphSpec used to describe a computational graph.

Each torch module has its own specification of input and output variables. The input specification i.e. args and kwargs of the torch modules forward function is represented by SampleMetadata. Those inputs can change computational graph. AITune treats each unique input specification as a separate graph which is tuned separately. This object represents such a computational graph with a name and input_spec information.

dynamic_shapes

DynamicShapes = field(default_factory=dict)

forward_signature

ForwardSignature

input_spec

SampleMetadata

name

str

output_spec

SampleMetadata

aitune.torch.module.graph_spec.GraphSpec.__repr__() -> str

Return representation of GraphSpec.

aitune.torch.module.graph_spec.GraphSpec.__str__() -> str

Return string representation of GraphSpec.

aitune.torch.module.graph_spec.GraphSpec._iter_batch_size_ranges(
    normalized: bool = False
) -> collections.abc.Iterator[tuple[int, int]]

Yield batch ranges from explicit definitions or inferred metadata.

aitune.torch.module.graph_spec.GraphSpec.from_dict(
    data: dict[str, typing.Any]
) -> aitune.torch.module.graph_spec.GraphSpec

staticmethod

Create GraphSpec from dictionary.

aitune.torch.module.graph_spec.GraphSpec.get_effective_input_shapes(
    locator: aitune.torch.module.locator.Locator,
    tensor_spec: aitune.torch.module.tensor_spec.TensorSpec
) -> tuple[list[int], list[int], list[int]]

Return the minimum, optimal, and maximum shapes to use for compilation.

An explicit user definition takes precedence. Otherwise, use the recorded TensorSpec bounds with its maximum shape as the optimal shape.

aitune.torch.module.graph_spec.GraphSpec.get_max_batch_size(
    normalized: bool = False
) -> int

Get max batch size from input spec.

Parameters:

normalized

boolDefaults to False

Flag to normalize the batch size against the global batch size.

aitune.torch.module.graph_spec.GraphSpec.get_min_batch_size() -> int | None

Get min batch size from input spec.

aitune.torch.module.graph_spec.GraphSpec.get_shape_definition(
    locator: aitune.torch.module.locator.Locator
) -> aitune.torch.dynamic_shapes.ShapeDefinition | None

Return the explicit shape definition for an input tensor, if configured.

aitune.torch.module.graph_spec.GraphSpec.make_batch(
    args: tuple,
    kwargs: dict[str, typing.Any],
    batch_size: int
) -> tuple[tuple, dict[str, typing.Any]]

Return a normalized call resized to the specified batch size.

aitune.torch.module.graph_spec.GraphSpec.to_dict() -> dict[str, typing.Any]

Convert the graph specification to a serializable dictionary.

aitune.torch.module.graph_spec.GraphSpec.update_max_batch_size(
    sample: tuple[tuple, dict],
    max_batch_size: int
) -> None

Update input metadata with the specified maximum batch size.

aitune.torch.module.graph_spec.GraphSpec.update_shapes_seen(
    inputs_metadata: aitune.torch.module.sample_metadata.SampleMetadata,
    outputs_metadata: aitune.torch.module.sample_metadata.SampleMetadata
)

Update input spec with other input spec.