Execution Graphs | NVIDIA AITune Documentation

Overview

During the tuning process, AITune analyzes module inputs to detect unique execution graphs—distinct computational paths through your model based on input characteristics. Understanding execution graphs is crucial because:

Separate Optimization: Each graph is tuned independently with its own optimized backend
Dynamic Shape Support: Graphs capture relationships between batch sizes, dynamic dimensions, and static shapes
Input Routing: At inference time, inputs are automatically routed to the correct optimized graph

What Defines an Execution Graph?

AITune creates a new execution graph when it encounters inputs that differ in these ways:

Tensor Rank: Tensors with different numbers of dimensions

1 module(torch.randn(1, 10))       # Graph 0: rank-2 tensor
2 module(torch.randn(1, 10, 5))    # Graph 1: rank-3 tensor (different graph!)

Argument Structure: Different supplied parameters or nested input structures

1 # Given forward(x, mask=False):
2 module(x, True)                 # Graph 0
3 module(x=x, mask=True)          # Graph 0: equivalent call
4 module(x)                       # Graph 1: mask is not supplied

Passing the same forward parameters positionally, by keyword, or with a mixture of both produces the same graph.

Non-Tensor Arguments (in strict mode): Different primitive values or configurations

1 module(x, mode="train")    # Graph 0
2 module(x, mode="eval")     # Graph 1 (if strict_mode=True)

Important: Tensors with the same rank but different shapes belong to the same graph. AITune handles shape variations through dynamic shape tracking (batch axes and dynamic dimensions).

Graph Detection in the Tuning Workflow

Execution graphs are detected during the Sample Gathering Phase of tuning:

Your model is executed with samples from the dataset
Each wrapped module records input/output metadata using SampleMetadata
AITune compares metadata to identify unique graph patterns
Each unique pattern becomes a separate GraphSpec
During Module Tuning, each graph is optimized independently

For a complete overview of the tuning process, see Tuning Workflow.

What This Guide Covers

This guide explores the technical details of execution graphs through the lens of SampleMetadata, the core class responsible for:

Capturing tensor metadata (shapes, dtypes, structure)
Detecting batch axes vs. dynamic dimensions
Tracking shape ranges (min/max) across samples
Enabling dynamic batching and shape inference

By the end of this guide, you’ll understand how AITune:

Identifies which inputs belong to the same graph
Learns dynamic shape patterns from multiple samples
Uses this information to optimize each execution path separately

Introduction to SampleMetadata

What is SampleMetadata?

SampleMetadata is a class designed to capture and track metadata about function inputs and outputs, particularly focusing on PyTorch tensors. It serves several important purposes:

Tensor Tracking: Automatically discovers and tracks all tensors in complex nested data structures
Shape Inference: Learns about dynamic dimensions and batch axes by observing multiple samples
Model Optimization: Enables optimization backends to understand input/output characteristics
Dynamic Batching: Supports scaling tensors to different batch sizes based on learned patterns

Key Concepts

Locators: Navigate through nested structures (tuples, lists, dicts, dataclasses and registered user types) to find tensors
TensorSpec: Underlying representation that tracks shape, dtype, and batch axis information
Dynamic Dimensions: Dimensions that vary across samples (e.g., sequence length in NLP)
Batch Axes: Dimensions that scale proportionally with batch size

Let’s start with a simple example:

1 from dataclasses import dataclass
2 import torch
3 from aitune.torch.module.sample_metadata import SampleMetadata, InfoLevel
4 
5 # Create a simple tensor and capture its metadata
6 simple_tensor = torch.randn(2, 3, 4)
7 inputs = {"simple_tensor": simple_tensor}
8 
9 metadata = SampleMetadata.from_inputs(inputs)
10 print(repr(metadata))

Tensors:
╒═══════════════╤═════════════════╤═══════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape     │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═══════════╪═════════════╪═════════════╪═══════════════╡
│ simple_tensor │ simple_tensor   │ [2, 3, 4] │ [2, 3, 4]   │ [2, 3, 4]   │ torch.float32 │
╘═══════════════╧═════════════════╧═══════════╧═════════════╧═════════════╧═══════════════╛

The output shows that SampleMetadata automatically detected the tensor and captured its forward parameter path, shape, and data type.

Creating Metadata from Inputs

The primary way to create SampleMetadata is through the from_inputs() static method. This method accepts:

inputs: A dictionary keyed by the module’s forward parameter names
strict: Boolean flag controlling whether to track non-tensor data (default: False)

AITune obtains this dictionary by binding a normalized (args, kwargs) call to its saved forward signature:

1 forward_inputs = forward_signature.normalize(args, kwargs)
2 metadata = SampleMetadata.from_inputs(forward_inputs.arguments)

Equivalent positional and keyword calls therefore produce the same parameter-keyed input representation.

Let’s explore different input patterns:

1 # Example 1: Multiple tensor parameters
2 inputs = {
3     "first": torch.randn(2, 3),
4     "second": torch.randn(4, 5),
5 }
6 
7 meta1 = SampleMetadata.from_inputs(inputs)
8 print("Example 1 - Multiple parameters:")
9 print(repr(meta1))

Example 1 - Multiple parameters:
Tensors:
╒═══════════════╤═════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ first         │ first           │ [2, 3]  │ [2, 3]      │ [2, 3]      │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ second        │ second          │ [4, 5]  │ [4, 5]      │ [4, 5]      │ torch.float32 │
╘═══════════════╧═════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛

1 # Example 2: Mixed primitives and tensors
2 inputs = {
3     "label": "some_string",       # Primitive - ignored by default
4     "tensor": torch.randn(2, 2),   # Tensor - tracked
5     "steps": 42,                   # Primitive - ignored by default
6     "data": torch.randn(3, 3),
7     "learning_rate": 0.001,        # Primitive - ignored by default
8 }
9 
10 meta2 = SampleMetadata.from_inputs(inputs)
11 print("Example 2 - Mixed types (strict=False):")
12 print(repr(meta2))
13 print("\nNotice that only tensors are tracked!")

Example 2 - Mixed types (strict=False):
Tensors:
╒═══════════════╤═════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ tensor        │ tensor          │ [2, 2]  │ [2, 2]      │ [2, 2]      │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ data          │ data            │ [3, 3]  │ [3, 3]      │ [3, 3]      │ torch.float32 │
╘═══════════════╧═════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛

Notice that only tensors are tracked!

Strict vs. Non-Strict Mode

By default during tuning, AITune operates in strict mode (config.strict_mode=True), which means SampleMetadata captures both tensors and non-tensor data (primitives, strings, etc.). This ensures different argument values create different execution graphs.

However, when calling SampleMetadata.from_inputs() directly with strict=False, it only tracks tensors and ignores all other data types. This is useful when you only care about tensor shapes for optimization purposes.

Strict mode (strict=True) is useful for:

Validating that function signatures match expected patterns
Debugging data flow through complex pipelines
Ensuring reproducibility of function calls

Let’s compare the two modes:

1 # Same inputs, different modes
2 inputs = {
3     "values": (1, 2, 3, torch.randn(2, 2)),
4     "t": torch.randn(2, 3),
5     "other": "abc",
6 }
7 
8 # Non-strict mode (default)
9 meta_non_strict = SampleMetadata.from_inputs(inputs, strict=False)
10 print("Non-Strict Mode (strict=False):")
11 print(repr(meta_non_strict))
12 print("\n" + "="*80 + "\n")
13 
14 # Strict mode
15 meta_strict = SampleMetadata.from_inputs(inputs, strict=True)
16 print("Strict Mode (strict=True):")
17 print(repr(meta_strict))

Non-Strict Mode (strict=False):
Tensors:
╒═══════════════╤═════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ values[3]     │ ('values', 3)   │ [2, 2]  │ [2, 2]      │ [2, 2]      │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ t             │ t               │ [2, 3]  │ [2, 3]      │ [2, 3]      │ torch.float32 │
╘═══════════════╧═════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛
================================================================================
Strict Mode (strict=True):
Tensors:
╒═══════════════╤═════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ values[3]     │ ('values', 3)   │ [2, 2]  │ [2, 2]      │ [2, 2]      │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ t             │ t               │ [2, 3]  │ [2, 3]      │ [2, 3]      │ torch.float32 │
╘═══════════════╧═════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛
Other:
╒═══════════════╤═════════════════╤═════════╕
│ Access Path   │ Semantic Path   │ Value   │
╞═══════════════╪═════════════════╪═════════╡
│ values[0]     │ ('values', 0)   │ 1       │
├───────────────┼─────────────────┼─────────┤
│ values[1]     │ ('values', 1)   │ 2       │
├───────────────┼─────────────────┼─────────┤
│ values[2]     │ ('values', 2)   │ 3       │
├───────────────┼─────────────────┼─────────┤
│ other         │ other           │ abc     │
╘═══════════════╧═════════════════╧═════════╛

Notice that in strict mode, we see an additional “Other” section that includes the primitive values (1, 2, 3, and “abc”).

Working with Nested Structures

One of the most powerful features of SampleMetadata is its ability to handle deeply nested data structures. Real-world model inputs often involve complex combinations of:

Tuples and Lists: For variable-length sequences
Dictionaries: For named parameters
Dataclasses: For structured configuration objects

SampleMetadata uses Locators to navigate these structures and find all tensors, no matter how deeply nested they are.

Let’s create a complex nested example:

1 # Define a custom dataclass
2 @dataclass
3 class ModelInput:
4     data: torch.Tensor
5     metadata: str
6 
7 # Create complex nested structure keyed by forward parameter name
8 inputs = {
9     "values": [
10         "first_arg",
11         torch.randn(1),                                      # Simple tensor
12         (torch.randn(2), torch.randn(3)),                    # Tuple of tensors
13         {"t": torch.randn(4)},                              # Dict with tensor
14         ModelInput(data=torch.randn(5), metadata="info"),   # Dataclass with tensor
15     ],
16     "t1": torch.randn(1, 1),
17     "t2": [torch.randn(2, 2), torch.randn(3, 3)],        # List of tensors
18     "t3": ModelInput(data=torch.randn(4, 4), metadata="xyz"),
19     "last": "other",
20 }
21 
22 nested_meta = SampleMetadata.from_inputs(inputs, strict=True)
23 print(repr(nested_meta))

Tensors:
╒════════════════╤═══════════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path    │ Semantic Path         │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞════════════════╪═══════════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ values[1]      │ ('values', 1)         │ [1]     │ [1]         │ [1]         │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ values[2][0]   │ ('values', 2, 0)      │ [2]     │ [2]         │ [2]         │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ values[2][1]   │ ('values', 2, 1)      │ [3]     │ [3]         │ [3]         │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ values[3]["t"] │ ('values', 3, 't')    │ [4]     │ [4]         │ [4]         │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ values[4].data │ ('values', 4, 'data') │ [5]     │ [5]         │ [5]         │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ t1             │ t1                    │ [1, 1]  │ [1, 1]      │ [1, 1]      │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ t2[0]          │ ('t2', 0)             │ [2, 2]  │ [2, 2]      │ [2, 2]      │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ t2[1]          │ ('t2', 1)             │ [3, 3]  │ [3, 3]      │ [3, 3]      │ torch.float32 │
├────────────────┼───────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ t3.data        │ ('t3', 'data')        │ [4, 4]  │ [4, 4]      │ [4, 4]      │ torch.float32 │
╘════════════════╧═══════════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛
Other:
╒════════════════════╤═══════════════════════════╤═══════════╕
│ Access Path        │ Semantic Path             │ Value     │
╞════════════════════╪═══════════════════════════╪═══════════╡
│ values[0]          │ ('values', 0)             │ first_arg │
├────────────────────┼───────────────────────────┼───────────┤
│ values[4].metadata │ ('values', 4, 'metadata') │ info      │
├────────────────────┼───────────────────────────┼───────────┤
│ t3.metadata        │ ('t3', 'metadata')        │ xyz       │
├────────────────────┼───────────────────────────┼───────────┤
│ last               │ last                      │ other     │
╘════════════════════╧═══════════════════════════╧═══════════╛

Understanding Paths

The Access Path column uses Python-like access syntax rooted at the real forward parameter name:

values[1]: Second element of the values parameter
values[2][0]: First element of the tuple at values[2]
values[3]["t"]: Value at key "t" in the dictionary at values[3]
values[4].data: The data attribute of the dataclass at values[4]
t2[0]: First element of the t2 parameter
t3.data: The data attribute of the t3 parameter

The Semantic Path column shows the stable identity used to match inputs and configure them, for example ("values", 3, "t"). Internally, a Locator also retains the container details required to retrieve or replace the value. Access paths are used in reports and as backend tensor IDs.

Describing Metadata - InfoLevel

SampleMetadata provides three levels of detail when displaying information, controlled by the InfoLevel enum:

InfoLevel.SHORT: Compact representation with just tensor access paths
InfoLevel.MEDIUM: Includes access paths, semantic paths, and current shapes (simple table format)
InfoLevel.FULL: Complete details including both paths, min/max shapes, and dtypes (fancy table format)

Let’s see the same metadata displayed at all three levels:

1 # Create sample metadata
2 inputs = {
3     "first": torch.randn(2, 3),
4     "second": torch.randn(4, 5, 6),
5     "mask": torch.randn(2, 1),
6 }
7 meta = SampleMetadata.from_inputs(inputs)
8 
9 print("InfoLevel.SHORT:")
10 print(meta.describe(InfoLevel.SHORT))
11 print("\n" + "="*80 + "\n")
12 
13 print("InfoLevel.MEDIUM:")
14 print(meta.describe(InfoLevel.MEDIUM))
15 print("\n" + "="*80 + "\n")
16 
17 print("InfoLevel.FULL:")
18 print(meta.describe(InfoLevel.FULL))

InfoLevel.SHORT:
Tensors: first, second, mask
================================================================================
InfoLevel.MEDIUM:
Tensors:
Access Path    Semantic Path    Shape
-------------  ---------------  ---------
first          first            [2, 3]
second         second           [4, 5, 6]
mask           mask             [2, 1]
================================================================================
InfoLevel.FULL:
Tensors:
╒═══════════════╤═════════════════╤═══════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape     │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═══════════╪═════════════╪═════════════╪═══════════════╡
│ first         │ first           │ [2, 3]    │ [2, 3]      │ [2, 3]      │ torch.float32 │
├───────────────┼─────────────────┼───────────┼─────────────┼─────────────┼───────────────┤
│ second        │ second          │ [4, 5, 6] │ [4, 5, 6]   │ [4, 5, 6]   │ torch.float32 │
├───────────────┼─────────────────┼───────────┼─────────────┼─────────────┼───────────────┤
│ mask          │ mask            │ [2, 1]    │ [2, 1]      │ [2, 1]      │ torch.float32 │
╘═══════════════╧═════════════════╧═══════════╧═════════════╧═════════════╧═══════════════╛

The FULL level is particularly useful because it shows:

Min Shape: The smallest dimensions seen for each axis
Max Shape: The largest dimensions seen for each axis
Dtype: The PyTorch data type of the tensor

These become interesting when we start tracking multiple samples with different shapes.

Dynamic Shape Tracking

One of the most sophisticated features of SampleMetadata is its ability to learn about dynamic dimensions and batch axes by observing multiple samples with different shapes.

How It Works

When you call update_shapes_seen() with metadata from a different sample:

Batch Axis Detection: If a dimension scales proportionally with batch size and the multiplier is an integer, it’s marked as a batch axis (e.g., batch0, batch1)
Dynamic Dimension Detection: If a dimension changes but not proportionally to batch size, it’s marked as a dynamic dimension (e.g., dim0, dim1)
Min/Max Tracking: The minimum and maximum values seen for each dimension are tracked

Let’s see this in action:

1 # Create initial metadata with batch size 1
2 inputs_initial = {
3     "first": torch.randn(1),
4     "second": torch.randn(2),
5     "third": torch.randn(5),
6     "data": torch.randn(1, 10),  # First dim batch, second dynamic
7 }
8 
9 meta_initial = SampleMetadata.from_inputs(inputs_initial, strict=False, batch_size=1)
10 print("Initial Metadata (batch_size=1):")
11 print(meta_initial.describe(InfoLevel.FULL))

Initial Metadata (batch_size=1):
Tensors:
╒═══════════════╤═════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ first         │ first           │ [1]     │ [1]         │ [1]         │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ second        │ second          │ [2]     │ [2]         │ [2]         │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ third         │ third           │ [5]     │ [5]         │ [5]         │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ data          │ data            │ [1, 10] │ [1, 10]     │ [1, 10]     │ torch.float32 │
╘═══════════════╧═════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛

1 # Create second metadata with different shapes and batch size 2
2 inputs_second = {
3     "first": torch.randn(2),   # Doubled (batch axis)
4     "second": torch.randn(5),  # Changed but not proportionally (dynamic)
5     "third": torch.randn(15),  # Changed but not proportionally (dynamic)
6     "data": torch.randn(2, 25),  # First dim doubled, second changed
7 }
8 
9 meta_second = SampleMetadata.from_inputs(inputs_second, strict=False, batch_size=2)
10 print("Second Metadata (batch_size=2):")
11 print(meta_second.describe(InfoLevel.FULL))

Second Metadata (batch_size=2):
Tensors:
╒═══════════════╤═════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ first         │ first           │ [2]     │ [2]         │ [2]         │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ second        │ second          │ [5]     │ [5]         │ [5]         │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ third         │ third           │ [15]    │ [15]        │ [15]        │ torch.float32 │
├───────────────┼─────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ data          │ data            │ [2, 25] │ [2, 25]     │ [2, 25]     │ torch.float32 │
╘═══════════════╧═════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛

1 # Update the initial metadata with information from the second sample
2 meta_initial.update_shapes_seen(meta_second)
3 print("Updated Metadata (after seeing both samples):")
4 print(meta_initial.describe(InfoLevel.FULL))

Updated Metadata (after seeing both samples):
Tensors:
╒═══════════════╤═════════════════╤════════════════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape              │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪════════════════════╪═════════════╪═════════════╪═══════════════╡
│ first         │ first           │ ['batch0']         │ [1]         │ [2]         │ torch.float32 │
├───────────────┼─────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ second        │ second          │ ['dim0']           │ [2]         │ [5]         │ torch.float32 │
├───────────────┼─────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ third         │ third           │ ['dim0']           │ [5]         │ [15]        │ torch.float32 │
├───────────────┼─────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ data          │ data            │ ['batch0', 'dim1'] │ [1, 10]     │ [2, 25]     │ torch.float32 │
╘═══════════════╧═════════════════╧════════════════════╧═════════════╧═════════════╧═══════════════╛

Understanding the Results

After updating, notice how the shapes have been transformed:

batch0: Dimensions that doubled when batch size doubled (1→2)
dim0, dim1: Dimensions that changed but not proportionally to batch size
Min/Max Shape: Now show the range of values observed

This information is crucial for:

Model compilation: Backends can create optimized graphs for dynamic shapes
Memory planning: Knowing the range helps allocate appropriate buffers
Validation: Ensuring new inputs fall within expected ranges

Batch Manipulation

Once SampleMetadata has learned about batch axes through update_shapes_seen(), it can use the make_batch() method to scale tensors to a target batch size.

How `make_batch()` Works

The method uses batch axis multipliers to determine how to scale each dimension:

Multiplier = 1: Standard batch axis, scales linearly with batch size
Multiplier > 1: Stacked batch axis (e.g., when inputs are vertically stacked)
Slicing: If current size > target, slice the tensor
Repeating: If current size < target, repeat the tensor

Let’s see this in action:

1 # First, create metadata and teach it about batch axes
2 inputs1 = {
3     "features": torch.randn(1, 5),
4     "stacked": torch.randn(2, 3),
5     "mask": torch.randn(1, 10),
6 }
7 
8 meta = SampleMetadata.from_inputs(inputs1, batch_size=1)
9 
10 # Second sample with batch size 2
11 inputs2 = {
12     "features": torch.randn(2, 5),
13     "stacked": torch.randn(4, 3),
14     "mask": torch.randn(2, 10),
15 }
16 
17 meta2 = SampleMetadata.from_inputs(inputs2, batch_size=2)
18 meta.update_shapes_seen(meta2)
19 
20 print("Learned Metadata:")
21 print(meta.describe(InfoLevel.FULL))

Learned Metadata:
Tensors:
╒═══════════════╤═════════════════╤════════════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path   │ Semantic Path   │ Shape          │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════╪═════════════════╪════════════════╪═════════════╪═════════════╪═══════════════╡
│ features      │ features        │ ['batch0', 5]  │ [1, 5]      │ [2, 5]      │ torch.float32 │
├───────────────┼─────────────────┼────────────────┼─────────────┼─────────────┼───────────────┤
│ stacked       │ stacked         │ ['batch0', 3]  │ [2, 3]      │ [4, 3]      │ torch.float32 │
├───────────────┼─────────────────┼────────────────┼─────────────┼─────────────┼───────────────┤
│ mask          │ mask            │ ['batch0', 10] │ [1, 10]     │ [2, 10]     │ torch.float32 │
╘═══════════════╧═════════════════╧════════════════╧═════════════╧═════════════╧═══════════════╛

1 # Now use make_batch to scale to a larger batch size
2 original_inputs = {
3     "features": torch.randn(1, 5),
4     "stacked": torch.randn(2, 3),
5     "mask": torch.randn(1, 10),
6 }
7 
8 print("Original shapes:")
9 print(f"  features: {original_inputs['features'].shape}")
10 print(f"  stacked: {original_inputs['stacked'].shape}")
11 print(f"  mask: {original_inputs['mask'].shape}")
12 print()
13 
14 # Scale to batch size 10
15 batched_inputs = meta.make_batch(original_inputs, batch_size=10)
16 
17 print("After make_batch(batch_size=10):")
18 print(f"  features: {batched_inputs['features'].shape}")
19 print(f"  stacked: {batched_inputs['stacked'].shape}")
20 print(f"  mask: {batched_inputs['mask'].shape}")

Original shapes:
  features: torch.Size([1, 5])
  stacked: torch.Size([2, 3])
  mask: torch.Size([1, 10])
After make_batch(batch_size=10):
  features: torch.Size([10, 5])
  stacked: torch.Size([20, 3])
  mask: torch.Size([10, 10])

Notice how:

The first dimensions (batch axes) scaled to match the target batch size of 10
stacked has a multiplier of 2, so it scaled to 20 (10 × 2)
Non-batch dimensions (like the 5, 3, 10) remained unchanged

TensorSpec Deep Dive

SampleMetadata is actually a container for multiple TensorSpec objects, where each TensorSpec represents one tensor in the input/output structure.

TensorSpec Attributes

shape: Current shape representation (may include symbolic dimensions)
min_shape: Minimum dimensions observed
max_shape: Maximum dimensions observed
dtype: PyTorch data type
_bs_multipliers: Internal batch size multipliers for each axis

Input identity is not part of TensorSpec. The paired Locator provides the parameter path, while TensorSpec contains only tensor properties.

Let’s inspect TensorSpec objects directly:

1 # Create metadata with dynamic shapes
2 inputs1 = {"features": torch.randn(1, 5), "data": torch.randn(1, 10)}
3 meta = SampleMetadata.from_inputs(inputs1, batch_size=1)
4 
5 inputs2 = {"features": torch.randn(2, 5), "data": torch.randn(2, 20)}
6 meta2 = SampleMetadata.from_inputs(inputs2, batch_size=2)
7 meta.update_shapes_seen(meta2)
8 
9 # Access individual TensorSpec objects
10 print("Individual TensorSpec objects:\n")
11 for locator, tensor_spec in meta.tensor_data:
12     print(f"Path: {locator.display_path}")
13     print(f"  Shape: {tensor_spec.shape}")
14     print(f"  Min Shape: {tensor_spec.min_shape}")
15     print(f"  Max Shape: {tensor_spec.max_shape}")
16     print(f"  Dtype: {tensor_spec.dtype}")
17     print(f"  Has batch axis: {tensor_spec.has_batch_axis()}")
18     print(f"  Has dynamic axis: {tensor_spec.has_dynamic_axis()}")
19     print(f"  Batch multipliers: {tensor_spec.get_batch_axis_multipliers()}")

Individual TensorSpec objects:
Path: features
  Shape: ['batch0', 5]
  Min Shape: [1, 5]
  Max Shape: [2, 5]
  Dtype: torch.float32
  Has batch axis: True
  Has dynamic axis: False
  Batch multipliers: {0: 1}
Path: data
  Shape: ['batch0', 'batch1']
  Min Shape: [1, 10]
  Max Shape: [2, 20]
  Dtype: torch.float32
  Has batch axis: True
  Has dynamic axis: False
  Batch multipliers: {0: 1, 1: 10}

Useful TensorSpec Methods

has_batch_axis(): Returns True if the tensor has at least one batch dimension
has_dynamic_axis(): Returns True if the tensor has at least one dynamic dimension
get_batch_axis_multipliers(): Returns a dict mapping axis index to its batch multiplier
matches(other): Checks if two TensorSpecs are compatible

These methods are used internally by SampleMetadata to perform operations like make_batch().

Practical Example - Complete Workflow

Let’s put everything together with a realistic scenario: profiling a model with variable-length sequences (like in NLP tasks).

Scenario

We have a language model that takes:

Input IDs with shape (batch_size, sequence_length)
Attention mask with shape (batch_size, sequence_length)
Position IDs with shape (batch_size, sequence_length)

We’ll profile it with different batch sizes and sequence lengths to learn the dynamic shapes.

1 # Simulate model profiling
2 @dataclass
3 class ModelInputs:
4     input_ids: torch.Tensor
5     attention_mask: torch.Tensor
6     position_ids: torch.Tensor
7 
8 # Sample 1: batch_size=1, seq_len=10
9 sample1_inputs = {
10     "inputs": ModelInputs(
11         input_ids=torch.randint(0, 1000, (1, 10)),
12         attention_mask=torch.ones(1, 10),
13         position_ids=torch.arange(10).unsqueeze(0),
14     )
15 }
16 
17 metadata = SampleMetadata.from_inputs(sample1_inputs, batch_size=1, strict=False)
18 print("After Sample 1 (batch=1, seq_len=10):")
19 print(metadata.describe(InfoLevel.FULL))

After Sample 1 (batch=1, seq_len=10):
Tensors:
╒═══════════════════════╤══════════════════════════════╤═════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path           │ Semantic Path                │ Shape   │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════════════╪══════════════════════════════╪═════════╪═════════════╪═════════════╪═══════════════╡
│ inputs.input_ids      │ ('inputs', 'input_ids')      │ [1, 10] │ [1, 10]     │ [1, 10]     │ torch.int64   │
├───────────────────────┼──────────────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ inputs.attention_mask │ ('inputs', 'attention_mask') │ [1, 10] │ [1, 10]     │ [1, 10]     │ torch.float32 │
├───────────────────────┼──────────────────────────────┼─────────┼─────────────┼─────────────┼───────────────┤
│ inputs.position_ids   │ ('inputs', 'position_ids')   │ [1, 10] │ [1, 10]     │ [1, 10]     │ torch.int64   │
╘═══════════════════════╧══════════════════════════════╧═════════╧═════════════╧═════════════╧═══════════════╛

1 # Sample 2: batch_size=2, seq_len=15
2 sample2_inputs = {
3     "inputs": ModelInputs(
4         input_ids=torch.randint(0, 1000, (2, 15)),
5         attention_mask=torch.ones(2, 15),
6         position_ids=torch.arange(15).unsqueeze(0).repeat(2, 1),
7     )
8 }
9 
10 metadata2 = SampleMetadata.from_inputs(sample2_inputs, batch_size=2, strict=False)
11 metadata.update_shapes_seen(metadata2)
12 
13 print("After Sample 2 (batch=2, seq_len=15):")
14 print(metadata.describe(InfoLevel.FULL))

After Sample 2 (batch=2, seq_len=15):
Tensors:
╒═══════════════════════╤══════════════════════════════╤════════════════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path           │ Semantic Path                │ Shape              │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════════════╪══════════════════════════════╪════════════════════╪═════════════╪═════════════╪═══════════════╡
│ inputs.input_ids      │ ('inputs', 'input_ids')      │ ['batch0', 'dim1'] │ [1, 10]     │ [2, 15]     │ torch.int64   │
├───────────────────────┼──────────────────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ inputs.attention_mask │ ('inputs', 'attention_mask') │ ['batch0', 'dim1'] │ [1, 10]     │ [2, 15]     │ torch.float32 │
├───────────────────────┼──────────────────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ inputs.position_ids   │ ('inputs', 'position_ids')   │ ['batch0', 'dim1'] │ [1, 10]     │ [2, 15]     │ torch.int64   │
╘═══════════════════════╧══════════════════════════════╧════════════════════╧═════════════╧═════════════╧═══════════════╛

1 # Sample 3: batch_size=4, seq_len=20
2 sample3_inputs = {
3     "inputs": ModelInputs(
4         input_ids=torch.randint(0, 1000, (4, 20)),
5         attention_mask=torch.ones(4, 20),
6         position_ids=torch.arange(20).unsqueeze(0).repeat(4, 1),
7     )
8 }
9 
10 metadata3 = SampleMetadata.from_inputs(sample3_inputs, batch_size=4, strict=False)
11 metadata.update_shapes_seen(metadata3)
12 
13 print("After Sample 3 (batch=4, seq_len=20):")
14 print(metadata.describe(InfoLevel.FULL))

After Sample 3 (batch=4, seq_len=20):
Tensors:
╒═══════════════════════╤══════════════════════════════╤════════════════════╤═════════════╤═════════════╤═══════════════╕
│ Access Path           │ Semantic Path                │ Shape              │ Min Shape   │ Max Shape   │ Dtype         │
╞═══════════════════════╪══════════════════════════════╪════════════════════╪═════════════╪═════════════╪═══════════════╡
│ inputs.input_ids      │ ('inputs', 'input_ids')      │ ['batch0', 'dim1'] │ [1, 10]     │ [4, 20]     │ torch.int64   │
├───────────────────────┼──────────────────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ inputs.attention_mask │ ('inputs', 'attention_mask') │ ['batch0', 'dim1'] │ [1, 10]     │ [4, 20]     │ torch.float32 │
├───────────────────────┼──────────────────────────────┼────────────────────┼─────────────┼─────────────┼───────────────┤
│ inputs.position_ids   │ ('inputs', 'position_ids')   │ ['batch0', 'dim1'] │ [1, 10]     │ [4, 20]     │ torch.int64   │
╘═══════════════════════╧══════════════════════════════╧════════════════════╧═════════════╧═════════════╧═══════════════╛

Analysis

After observing three samples with different batch sizes and sequence lengths:

First dimension: Identified as batch0 because it scaled proportionally (1→2→4)
Second dimension: Identified as dim1 because it varied dynamically (10→15→20)
Min/Max ranges: Captured the observed ranges for both dimensions

This information can now be used by optimization backends to compile efficient code for these dynamic shapes.

1 # Now we can create inputs for any batch size!
2 test_inputs = {
3     "inputs": ModelInputs(
4         input_ids=torch.randint(0, 1000, (2, 12)),
5         attention_mask=torch.ones(2, 12),
6         position_ids=torch.arange(12).unsqueeze(0).repeat(2, 1),
7     )
8 }
9 
10 print("Original test input shapes:")
11 print(f"  input_ids: {test_inputs['inputs'].input_ids.shape}")
12 print(f"  attention_mask: {test_inputs['inputs'].attention_mask.shape}")
13 print(f"  position_ids: {test_inputs['inputs'].position_ids.shape}")
14 print()
15 
16 # Scale to batch size 8
17 scaled_inputs = metadata.make_batch(test_inputs, batch_size=8)
18 
19 print("After scaling to batch_size=8:")
20 print(f"  input_ids: {scaled_inputs['inputs'].input_ids.shape}")
21 print(f"  attention_mask: {scaled_inputs['inputs'].attention_mask.shape}")
22 print(f"  position_ids: {scaled_inputs['inputs'].position_ids.shape}")
23 print("\nNote: Batch dimension scaled to 8, but sequence length (dim1) remained at 12")

Original test input shapes:
  input_ids: torch.Size([2, 12])
  attention_mask: torch.Size([2, 12])
  position_ids: torch.Size([2, 12])
After scaling to batch_size=8:
  input_ids: torch.Size([8, 12])
  attention_mask: torch.Size([8, 12])
  position_ids: torch.Size([8, 12])
Note: Batch dimension scaled to 8, but sequence length (dim1) remained at 12

Summary

Key Takeaways

Purpose: SampleMetadata captures and tracks metadata about tensors in complex data structures, enabling model optimization and dynamic batching.
Creation: Use SampleMetadata.from_inputs(inputs, strict=bool) with inputs keyed by forward parameter name.
Strict Mode: Controls whether only tensors (strict=False) or all data types (strict=True) are tracked.
Nested Structures: Automatically handles tuples, lists, dicts, and dataclasses using Locators.
InfoLevel: Three display modes (SHORT, MEDIUM, FULL) provide different levels of detail.
Dynamic Shape Learning: update_shapes_seen() learns about batch axes and dynamic dimensions by observing multiple samples.
Batch Manipulation: make_batch() can scale tensors to any batch size based on learned batch axis multipliers.
TensorSpec: The underlying representation of each tensor, containing shape, dtype, and batch information.

Use in AI-Tune Pipeline

SampleMetadata is a fundamental building block in the AITune library, used by:

RecordingModule: Captures input/output metadata during profiling
Backends: Use metadata to configure optimized execution (TensorRT, TorchScript, etc.)
Graph Compilation: Enables the creation of optimized graphs for dynamic shapes

Source Code

For more details, see:

aitune/torch/module/sample_metadata.py
aitune/torch/module/tensor_spec.py
aitune/torch/module/locator.py

If you would like to tinker with SampleMetadata, you can find this example in notebooks/sample_metadata_walkthrough.ipynb.