> For clean Markdown content of this page, append .md to this URL.

# aitune.torch.module.sample_metadata

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

## Module Contents

### Classes

| Name                                                                    | Description                                                                |
| ----------------------------------------------------------------------- | -------------------------------------------------------------------------- |
| [`SampleMetadata`](#aitune-torch-module-sample_metadata-SampleMetadata) | Metadata description of inputs (args and kwargs) or outputs of a function. |

### Functions

| Name                                                                | Description                                                                      |
| ------------------------------------------------------------------- | -------------------------------------------------------------------------------- |
| [`batch_tensor`](#aitune-torch-module-sample_metadata-batch_tensor) | Batch tensor so that instead of its own batch size it uses specified batch size. |

### API

```python
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:

1. The **init** method should not be used directly. instead:
   * for inputs - use SampleMetadata.from\_inputs(bound\_arguments)
   * for outputs - use SampleMetadata.from\_outputs(output)

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

---

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

Equality operator.

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

Compute hash of sample metadata.

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

Return representation of metadata.

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

Convert sample metadata to string.

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

Get information describing sample metadata.

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

Check if dynamic axes are detected in the metadata.

```python
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

```python
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.

```python
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.

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

Check if metadata has batch axis.

```python
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

```python
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.

```python
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`.

```python
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.

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

Update shapes seen from other SampleMetadata.

```python
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