PyTorch: torch.export and explicit dynamic shape contracts

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: torch.export and explicit dynamic shape contracts

Concept and contract

Distinguish an exported program’s input guards from a normal eager call.

Modern PyTorch offers torch.export for capturing a program together with input assumptions and constraints.

That exported representation is different from simply saving weights or calling a Python module eagerly. Static example dimensions can become guards, so an exported program may reject a different batch shape even when the original eager model could compute it.

Mechanism and worked example

The fixture exports a known two-feature linear model using a two-example input and reconstructs its callable exported module.

The fixture exports a known two-feature linear model using a two-example input and reconstructs its callable exported module. The initial export is intentionally static.

A two-example call succeeds with known outputs, while a three-example call violates its batch constraint. The helper reports the actual guarded execution result rather than inferring acceptance from the eager model’s mathematical shape compatibility.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def exported_batch(batch_size):
    model = nn.Linear(2, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(1)
    example = torch.ones(2, 2)
    exported = torch.export.export(model, (example,))
    module = exported.module()
    try:
        output = module(torch.ones(batch_size, 2))
    except (RuntimeError, AssertionError):
        return {'accepted': False}
    return {'accepted': True, 'shape': list(output.shape), 'values': output.tolist()}

Reasoning: The fixture exports a known two-feature linear model using a two-example input and reconstructs its callable exported module. The initial export is intentionally static.

Counterexamples and common mistakes

The counterexample bypasses exported.module and calls the original eager model.

The counterexample bypasses exported.module and calls the original eager model. That accepts an otherwise forbidden batch and conceals the missing deployment contract.

Export is not a guarantee of universal hardware support or an automatic performance improvement. Backends, supported operations and dynamic constraints still need validation, and torch.compile has a different purpose and runtime tradeoff.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def exported_batch(batch_size):
    model = nn.Linear(2, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(1)
    example = torch.ones(2, 2)
    exported = torch.export.export(model, (example,))
    module = model
    try:
        output = module(torch.ones(batch_size, 2))
    except (RuntimeError, AssertionError):
        return {'accepted': False}
    return {'accepted': True, 'shape': list(output.shape), 'values': output.tolist()}

Reasoning: The counterexample bypasses exported.module and calls the original eager model. That accepts an otherwise forbidden batch and conceals the missing deployment contract.

Independent practice: one bounded framework change

Declare the first input dimension dynamic with a named Dim bounded from one through four.

Declare the first input dimension dynamic with a named Dim bounded from one through four. Keep the two-feature dimension static.

The exported module should then accept batches two and three but reject five. This is a small input-contract extension, not a compiler benchmark. Use the pinned version’s public dynamic_shapes API and preserve actual exported execution in the helper.

Step 1: Declare a dynamic batch dimension bounded from one through four. Step 2: Pass dynamic_shapes to torch.export.export.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def exported_batch(batch_size):
    # Remove this placeholder, then adapt the retained baseline:
    # Declare a dynamic batch dimension bounded from one through four.
    raise NotImplementedError
    model = nn.Linear(2, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(1)
    example = torch.ones(2, 2)
    exported = torch.export.export(model, (example,))
    module = exported.module()
    try:
        output = module(torch.ones(batch_size, 2))
    except (RuntimeError, AssertionError):
        return {'accepted': False}
    return {'accepted': True, 'shape': list(output.shape), 'values': output.tolist()}

Reasoning: Declare a dynamic batch dimension bounded from one through four. Pass dynamic_shapes to torch.export.export.

Hint 1: Keep feature dimension two static.
Hint 2: Run the exported callable and test an out-of-bound batch.

Solution, tests and trade-offs

Baseline cases compare batch two, three and one under the static export.

Baseline cases compare batch two, three and one under the static export. Practice verifies two and three are accepted while five is rejected.

Exact output shapes and known values demonstrate that the exported callable runs. The fixture uses CPU and a tiny linear model; compilation speed, deployment backend coverage and production latency are deliberately not inferred from these checks.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

from torch import nn

def exported_batch(batch_size):
    model = nn.Linear(2, 1, bias=False)
    with torch.no_grad():
        model.weight.fill_(1)
    example = torch.ones(2, 2)
    exported = torch.export.export(model, (example,), dynamic_shapes=({0: torch.export.Dim('batch', min=1, max=4)},))
    module = exported.module()
    try:
        output = module(torch.ones(batch_size, 2))
    except (RuntimeError, AssertionError):
        return {'accepted': False}
    return {'accepted': True, 'shape': list(output.shape), 'values': output.tolist()}

Reasoning: Baseline cases compare batch two, three and one under the static export. Practice verifies two and three are accepted while five is rejected.

Interview reasoning and next step

Choose an artifact according to its consumer: state_dict for known Python architectures, export for an explicitly constrained program flow, and an appropriate backend-specific process for deployment..

Choose an artifact according to its consumer: state_dict for known Python architectures, export for an explicitly constrained program flow, and an appropriate backend-specific process for deployment.

Document allowed shapes, dtypes and devices alongside version compatibility. Prefer current export guidance over treating older TorchScript examples as the default modern path, while maintaining existing systems according to their actual supported toolchain.

Optional stretch — after the core task works: Document dtype and shape guards for a separately versioned inference artifact.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-export. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Distinguish an exported program’s input guards from a normal eager call.