aitune.records.shapes

View as Markdown

Frontend-neutral tensor contracts and tuning metadata.

Module Contents

Classes

NameDescription
TensorSpecDescribe one tensor in an executable interface.
TunedTensorSpecDescribe one tensor over the shape range AITune selected for an artifact.

API

class aitune.records.shapes.TensorSpec(
name: str,
shape: tuple[int | str | None, ...]
)
Dataclass

Describe one tensor in an executable interface.

Integer dimensions are fixed, strings are symbolic dimensions, and None is an unnamed dynamic dimension. The specification contains only information declared by the executable; concrete bounds belong to :class:TunedTensorSpec.

Parameters:

name
str

Tensor name used by the executable.

dtype
DType

Tensor element type.

shape
tuple[int | str | None, ...]

Static and dynamic dimensions declared by the executable.

dtype
DType
name
str
shape
tuple[int | str | None, ...]
aitune.records.shapes.TensorSpec.__post_init__() -> None

Validate the tensor name and declared dimensions.

class aitune.records.shapes.TunedTensorSpec(
name: str,
shape: tuple[int | str | None, ...],
min_shape: tuple[int, ...],
max_shape: tuple[int, ...],
batch_axis: int | None = None
)
Dataclass

Bases: TensorSpec

Describe one tensor over the shape range AITune selected for an artifact.

batch_axis is None when tuning did not identify a logical request batch axis.

Parameters:

name
str

Tensor name used by the executable.

dtype
DType

Tensor element type.

shape
tuple[int | str | None, ...]

Static and dynamic dimensions declared by the executable.

min_shape
tuple[int, ...]

Smallest shape AITune selected for the artifact.

max_shape
tuple[int, ...]

Largest shape AITune selected for the artifact.

batch_axis
int | NoneDefaults to None

Logical batch-axis index, or None when unknown.

batch_axis
int | None = None
max_batch_size
int | None

Return the largest batch size AITune selected, if known.

max_shape
tuple[int, ...]
min_batch_size
int | None

Return the smallest batch size AITune selected, if known.

min_shape
tuple[int, ...]
aitune.records.shapes.TunedTensorSpec.__post_init__() -> None

Validate concrete bounds against the static specification.

min_shape: tuple[int, ...],
max_shape: tuple[int, ...],
batch_axis: int | None = None
classmethod

Add tuning results to an existing tensor specification.

Parameters:

spec
TensorSpec

Static tensor specification.

min_shape
tuple[int, ...]

Smallest shape AITune selected for the artifact.

max_shape
tuple[int, ...]

Largest shape AITune selected for the artifact.

batch_axis
int | NoneDefaults to None

Logical batch-axis index, or None when unknown.

Returns: TunedTensorSpec

A tuned tensor specification.