aitune.torch.module.sample_metadata | NVIDIA AITune Documentation

Contains SampleMetadata which represents metadata of a function inputs (args and kwargs) or outputs.

Module Contents

Classes

Name	Description
`SampleMetadata`	Metadata description of inputs (args and kwargs) or outputs of a function.

Functions

Name	Description
`batch_tensor`	Batch tensor so that instead of its own batch size it uses specified batch size.

API

class aitune.torch.module.sample_metadata.SampleMetadata(
    tensor_data: tuple[tuple[aitune.torch.module.locator.Locator, aitune.torch.module.tensor_spec.TensorSpec], ...],
    other_data: tuple[tuple[aitune.torch.module.locator.Locator, typing.Any], ...],
    strict: bool = False,
    llm_phase: typing.Literal['prefill', 'decode'] | None = None
)

Metadata description of inputs (args and kwargs) or outputs of a function.

SampleMetadata captures and tracks metadata about function inputs (args and kwargs) and outputs. It serves several important purposes:

Tensor Tracking: Automatically discovers and tracks all tensors in nested structures (tuples, lists, dicts, dataclasses, custom objects).
Shape Inference: Learns about dynamic dimensions and batch axes by observing multiple samples with different shapes via update_shapes_seen().
Dynamic Batching: Supports scaling tensors to different batch sizes based on learned patterns using make_batch().

Note:

The init method should not be used directly. instead:
- for inputs - use SampleMetadata.from_inputs(bound_arguments)
- for outputs - use SampleMetadata.from_outputs(output)
In order to support Hugging Face integrations with kv cache the following behavior is changed:
- the equality and hash operator ignore all tensors and return value corresponding to the llm_graph_type, which is either “prefill” or “decode”. This is deduced based on the cache_position tensor. It cannot be done otherwise because even though kv cache is lazily initialized, it can be cached on subsequent generate calls i.e. it can always have tensors irrespective of the phase.

llm_phase

Literal['prefill', 'decode', '']

Get LLM graph type.

other_data

tuple[tuple[Locator, Any], ...]

Get list of other data.

tensor_data

tuple[tuple[Locator, TensorSpec]]

Get list of tensor data.

tensor_specs

list[TensorSpec]

Get list of tensor specs.

aitune.torch.module.sample_metadata.SampleMetadata.__eq__(
    __value: object
) -> bool

Equality operator.

aitune.torch.module.sample_metadata.SampleMetadata.__hash__() -> int

Compute hash of sample metadata.

aitune.torch.module.sample_metadata.SampleMetadata.__repr__()

Return representation of metadata.

aitune.torch.module.sample_metadata.SampleMetadata.__str__() -> str

Convert sample metadata to string.

aitune.torch.module.sample_metadata.SampleMetadata.describe(
    info_level: aitune.torch.module.tensor_spec.InfoLevel = InfoLevel.FULL
) -> str

Get information describing sample metadata.

aitune.torch.module.sample_metadata.SampleMetadata.detected_dynamic_axis() -> bool

Check if dynamic axes are detected in the metadata.

aitune.torch.module.sample_metadata.SampleMetadata.from_dict(
    data: dict
) -> aitune.torch.module.sample_metadata.SampleMetadata

staticmethod

Create SampleMetadata from dictionary.

Parameters:

data

dict

Dictionary containing serialized metadata

Returns: SampleMetadata

A SampleMetadata instance

aitune.torch.module.sample_metadata.SampleMetadata.from_inputs(
    inputs: dict[str, typing.Any],
    strict: bool = False,
    batch_size: int | None = None
) -> aitune.torch.module.sample_metadata.SampleMetadata

staticmethod

Create SampleMetadata from inputs keyed by forward parameter name.

If strict is True, then other data is also included.

aitune.torch.module.sample_metadata.SampleMetadata.from_outputs(
    output: typing.Any,
    strict: bool = False,
    batch_size: int | None = None
) -> aitune.torch.module.sample_metadata.SampleMetadata

staticmethod

Create SampleMetadata from outputs.

If strict is True, then other data is also included.

aitune.torch.module.sample_metadata.SampleMetadata.has_batch_axis() -> bool

Check if metadata has batch axis.

aitune.torch.module.sample_metadata.SampleMetadata.make_batch(
    inputs: dict[str, typing.Any],
    batch_size: int
) -> dict[str, typing.Any]

Extrapolate all tensors in forward inputs to the specified batch size.

Parameters:

inputs

dict[str, Any]

Inputs keyed by forward parameter name

batch_size

int

Batch size

Returns: dict[str, Any]

Bound inputs with all tensors having the specified batch size

aitune.torch.module.sample_metadata.SampleMetadata.to_dict() -> dict

Convert sample metadata to a serializable dictionary.

Returns: dict

A dictionary representation of the metadata that can be serialized.

aitune.torch.module.sample_metadata.SampleMetadata.to_json_dict() -> dict

Convert sample metadata to a JSON-serializable dictionary.

Returns: dict

A dictionary representation safe for json.dumps.

aitune.torch.module.sample_metadata.SampleMetadata.update_max_batch_size(
    inputs: dict[str, typing.Any],
    max_batch_size: int
)

Update input spec with max batch size information.

aitune.torch.module.sample_metadata.SampleMetadata.update_shapes_seen(
    other: aitune.torch.module.sample_metadata.SampleMetadata
)

Update shapes seen from other SampleMetadata.

aitune.torch.module.sample_metadata.batch_tensor(
    tensor: torch.Tensor,
    tensor_spec: aitune.torch.module.tensor_spec.TensorSpec,
    batch_size: int
) -> torch.Tensor

Batch tensor so that instead of its own batch size it uses specified batch size.

Parameters:

tensor

torch.Tensor

Tensor to batch

tensor_spec

TensorSpec

Metadata to batch tensor according to

batch_size

int

Batch size

Returns: torch.Tensor

Tensor of specified batch size