aitune.torch.module.tensor_spec | NVIDIA AITune Documentation

Contains TensorSpec which represents metadata of a tensor.

Module Contents

Classes

Name	Description
`InfoLevel`	Enum representing different levels of information detail.
`TensorSpec`	TensorSpec is used to describe tensor metadata.

API

class aitune.torch.module.tensor_spec.InfoLevel

Bases: enum.Enum

Enum representing different levels of information detail.

FULL

= auto()

MEDIUM

= auto()

SHORT

= auto()

class aitune.torch.module.tensor_spec.TensorSpec(
    shape: list[str | int],
    min_shape: list[int],
    max_shape: list[int],
    dtype: torch.dtype | None,
    _bs_multipliers: list[float]
)

Dataclass

TensorSpec is used to describe tensor metadata.

_bs_multipliers

list[float]

dtype

dtype | None

max_shape

list[int]

min_shape

list[int]

shape

list[str | int]

aitune.torch.module.tensor_spec.TensorSpec.__eq__(
    other: object
) -> bool

Check if two TensorSpec are equal.

Tensors of the same rank are considered equal. Particular dimensions sizes can be different as there can be dynamic dimensions.

aitune.torch.module.tensor_spec.TensorSpec.__hash__() -> int

Hash of TensorSpec.

Tensors of the same rank are considered equal. Particular dimensions sizes can be different as there can be dynamic dimensions.

aitune.torch.module.tensor_spec.TensorSpec.__repr__() -> str

Get representation of TensorSpec.

aitune.torch.module.tensor_spec.TensorSpec.__str__() -> str

Get string representation of TensorSpec.

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

Get information describing TensorSpec.

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

staticmethod

Create TensorSpec from dictionary.

aitune.torch.module.tensor_spec.TensorSpec.from_tensor(
    tensor: torch.Tensor,
    batch_size: int
)

staticmethod

Create TensorSpec from tensor.

Parameters:

tensor

torch.Tensor

Tensor to create TensorSpec from

batch_size

int

Batch size

aitune.torch.module.tensor_spec.TensorSpec.get_batch_axis_multipliers() -> dict[int, int]

Return mapping for batch axis and its multiplier.

aitune.torch.module.tensor_spec.TensorSpec.get_max_batch_size() -> int

Get max batch size from tensor spec.

aitune.torch.module.tensor_spec.TensorSpec.get_min_batch_size() -> int | None

Get min batch size from tensor spec.

aitune.torch.module.tensor_spec.TensorSpec.has_batch_axis() -> bool

Check if tensor has batch axis.

aitune.torch.module.tensor_spec.TensorSpec.has_dynamic_axis() -> bool

Check if tensor has dynamic axis.

aitune.torch.module.tensor_spec.TensorSpec.is_int(
    value: float
) -> bool

staticmethod

Check if value is an integer by checking value not the type.

aitune.torch.module.tensor_spec.TensorSpec.matches(
    other: aitune.torch.module.tensor_spec.TensorSpec
) -> bool

Check if tensor spec matches other tensor spec.

Tensor spec matches other if ranks are same and each ordinal dimension is the same.

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

Convert TensorSpec to a serializable dictionary.

aitune.torch.module.tensor_spec.TensorSpec.update_shapes_seen(
    other: aitune.torch.module.tensor_spec.TensorSpec
)

Update shapes seen from other TensorSpec.

Tensor have to have same rank in order to update self.

The algorithm for detecting batch dimension is the following: given two different batch sizes, if batch size multiplier is the same for same axis in both tensors and is an integer, then it is a batch dimension otherwise it is a dynamic dimension.

The reason for using multipliers is that some models stack input tensor vertically and the resulting input has double of the batch size i.e. local batch size is 2x the global batch size.

Example of the algorithm - let’s assume we observed tensor[1, 2, 3, 4] given bs=1. We calculated multipliers to be [1, 2, 3, 4]. Now if we see tensor[2, 8, 6, 4] with bs=2, we calculate multipliers to be [1, 4, 3, 2]. We can make some conclusions:

0th axis is batch axis, multiplier is 1
1st axis is dynamic axis, multiplier is 2 and 4 - this could be for example length in LLMs
2nd axis is batch axis, multiplier is 3 - the input tensor is v-stacked thus multiplier is 3
3rd axis is static axis - never changes w.r.t. batch size

This algorithm is not foolproof and can fail in some cases e.g. sequence length in LLMs is equal to batch size which is unlikely to happen in practice. However to mitigate this risk, there is additional check for batch axis multiplier which has to be an integer.