Symbolic Shapes in torch.compile

Table of Contents

• 1. Motivation: The Static vs. Dynamic Shape Tradeoff
  • 1.1 Static Shapes and Cache Thrashing
  • 1.2 Dynamic Shapes and Their Costs
  • 1.3 The Overhead Budget
• 2. ShapeEnv: The Central State Manager
  • 2.1 Core Data Structures
  • 2.2 Entry Point: create_symbolic_sizes_strides_storage_offset
  • 2.3 SymNode: The Internal Carrier
• 3. Symbolic Integers: SymPy Expressions Throughout the Stack
  • 3.1 From Symbol to Expression Tree
  • 3.2 SymPy Assumptions and Simplification
  • 3.3 User-Visible Types: SymInt, SymFloat, SymBool
• 4. Guard Generation
  • 4.1 What Triggers a Guard
  • 4.2 Guard Serialization and Source Annotations
  • 4.3 Guard Types and Their Costs
  • 4.4 Divisibility Guards
• 5. Specialization on 0 and 1
  • 5.1 Why the Special Cases Exist
  • 5.2 Performance Implication
  • 5.3 Disabling Zero-One Specialization
• 6. Backed vs. Unbacked SymInts
  • 6.1 Backed SymInts
  • 6.2 Unbacked SymInts
  • 6.3 Deferred Runtime Asserts
  • 6.4 torch._check and Range Refinement
• 7. The mark_dynamic / mark_static / torch.export.Dim APIs
  • 7.1 mark_dynamic
  • 7.2 mark_static
  • 7.3 torch.export.Dim
  • 7.4 How Annotations Feed into ShapeEnv
• 8. Symbolic Shapes in TorchInductor
  • 8.1 Loop IR with Symbolic Bounds
  • 8.2 Triton Kernel Parameterization
  • 8.3 Constant Folding for Static Dimensions
  • 8.4 Guard Tightening in Inductor
• 9. Automatic Dynamic Shapes: dynamic=“auto”
• 10. Debugging Symbolic Shapes
  • 10.1 TORCH_LOGS=“+dynamic”
  • 10.2 torch._dynamo.explain
  • 10.3 ShapeEnv as a Test Harness
  • 10.4 Common Failure Modes
• References

1. Motivation: The Static vs. Dynamic Shape Tradeoff

1.1 Static Shapes and Cache Thrashing

When TorchDynamo traces a function with a concrete input, every tensor dimension is a known integer. The compiled artifact — FX graph, fused Triton kernels, loop bounds, tile size decisions — is baked for those exact dimensions.

Definition (Shape Specialization). A compiled artifact \(A\) is shape-specialized if for every tensor dimension \(d_i\) appearing in the original function, \(A\) contains the concrete value \(d_i = c_i\) for some \(c_i \in \mathbb{Z}_{>0}\) (or \(c_i = 0\)) as a compile-time constant rather than a runtime parameter.

The benefits are substantial:

• Loop bounds are constants. A loop for i in range(s0): with \(s_0 = 64\) becomes a constant-bounded loop; the compiler can unroll, pipeline, or choose tile sizes exactly.
• Constant folding eliminates index arithmetic. Expressions like stride[0] = s1 * s2 with \(s_1 = 128\), \(s_2 = 4\) fold to 512 at compile time.
• Bounds checks are statically proved. Access patterns can be verified safe at compilation time.

The liability is equally severe: each distinct shape tuple triggers a fresh compilation. A transformer model serving variable-length sequences — lengths 64, 128, 192, 256, 512, … — produces one compiled artifact per length, each potentially taking seconds to compile. This cache thrashing is prohibitive for production inference.

1.2 Dynamic Shapes and Their Costs

Dynamic shape compilation replaces selected integer dimensions with sympy.Symbol objects — symbolic unknowns that flow through the entire compilation pipeline as algebraic expressions. The compiled artifact becomes a family of programs parameterized by the symbol values, valid for all conforming inputs without recompilation.

Definition (Dynamic Compilation). Let \(\mathbf{s} = (s_0, s_1, \ldots, s_{k-1})\) be a tuple of \(k\) free symbols in \(\mathbb{Z}_{>0}\). A compiled artifact \(A(\mathbf{s})\) is dynamically compiled with respect to \(\mathbf{s}\) if:

1. All loop bounds, strides, and index expressions in \(A\) are SymPy expressions in \(\mathbf{s}\).
2. A guard predicate \(G(\mathbf{s})\) — a conjunction of SymPy inequalities — is attached to \(A\); at runtime, \(G\) is evaluated with the concrete symbol values to certify that \(A\) is valid for the incoming shapes.

The costs of dynamic compilation are:

1. Symbolic arithmetic overhead at trace time. SymPy must simplify expressions instead of performing constant arithmetic. Trace time can be 2–4x longer.
2. Disabled static optimizations. Loop unrolling and tile size selection that depend on knowing an exact dimension cannot fire.
3. Runtime kernel arguments. Triton kernels receive \(s_0, s_1, \ldots\) as integer arguments on every launch. This adds a small kernel launch overhead (typically one mov instruction per symbol into the argument buffer).
4. Guard evaluation on each call. The guard predicate \(G(\mathbf{s})\) is evaluated as a small Python expression on each invocation.

1.3 The Overhead Budget

Surprisingly, the practical overhead of dynamic shapes versus static shapes is often only 5–10%. LLVM and the Triton compiler can constant-fold symbol values once they are known at kernel-call time, recovering many static optimizations at launch. The dominant cost is guard evaluation, which for typical models amounts to a few microseconds — negligible against GPU kernel latency of milliseconds.

Key principle: use dynamic shapes for dimensions that vary per request (batch size, sequence length); use static shapes for dimensions that are fixed by the model architecture (embedding dimension, number of heads).

2. ShapeEnv: The Central State Manager

2.1 Core Data Structures

Definition (ShapeEnv). torch.fx.experimental.symbolic_shapes.ShapeEnv is the singleton object allocated once per Dynamo compilation frame. It acts as the oracle for all symbolic shape decisions during that compilation and accumulates all constraints that must hold for the compiled artifact to remain valid.

The key internal data structures are:

A ShapeGuard is a small dataclass:

@dataclass
class ShapeGuard:
    expr: sympy.Expr      # a SymPy boolean expression
    stack: StackSummary   # traceback to where the guard was added

The ValueRanges type wraps a (lower, upper) pair of SymPy expressions (possibly \(-\infty\) / \(+\infty\)) and supports interval arithmetic: ValueRanges.add, ValueRanges.mul, etc.

2.2 Entry Point: create_symbolic_sizes_strides_storage_offset

The primary entry point for creating symbolic metadata for a tensor is:

sizes, strides, storage_offset = shape_env.create_symbolic_sizes_strides_storage_offset(
    t,       # a concrete torch.Tensor (or FakeTensor with a concrete shape)
    source,  # a Source object (e.g. LocalSource("x")) identifying the tensor
)

For each dimension \(i\):

1. If the dimension is marked static (via mark_static, or if the policy is dynamic=False, or if the value is 0 or 1 and zero-one specialization is active), sizes[i] is the concrete integer t.shape[i].
2. If the dimension is marked dynamic (via mark_dynamic, or if the policy is dynamic=True), a fresh sympy.Symbol \(s_k\) is allocated (with \(k\) incremented globally), \(s_k\) is inserted into var_to_val with value t.shape[i], its range is initialized in var_to_range to [2, inf] (size-like, positive, avoiding 0 and 1 per the specialization policy), and sizes[i] is a SymInt wrapping a SymNode that points to \(s_k\).

The same logic applies to strides (which can be symbolic when views are involved) and the storage offset.

2.3 SymNode: The Internal Carrier

Every SymInt visible in Python wraps a SymNode object. The SymNode holds:

• A SymPy expression (the symbolic value).
• A reference back to the owning ShapeEnv.
• Optional metadata: a hint (the concrete value from var_to_val), used to guide simplification without committing to a guard.

The entire symbolic reasoning apparatus lives in Python, not C++. This was a deliberate design choice: SymPy, traceback capture, and the guard serializer all require the Python runtime. The only C++ involvement is the dispatch layer that routes shape queries to SymNode when a tensor has symbolic metadata.

from torch.fx.experimental.symbolic_shapes import ShapeEnv, ValueRanges
import sympy

shape_env = ShapeEnv()
s0 = shape_env.create_symbol(64, source=..., dynamic_dim=True,
                              constraint_dim=ValueRanges(1, 512))
expr = s0**2 + 1
val = shape_env.evaluate_expr(expr)
# val == 4097 (the hint-evaluated result)
# Guard added: s0**2 + 1 == 4097  (an EQUALS_MATCH on the expression)

3. Symbolic Integers: SymPy Expressions Throughout the Stack

3.1 From Symbol to Expression Tree

Each dynamic dimension becomes a leaf sympy.Symbol, named s0, s1, …, sN for backed symbols and u0, u1, … for unbacked symbols (see §6). All arithmetic on SymInt objects in Python is overloaded to produce SymPy expression trees:

\[ \texttt{SymInt}(s_0) + \texttt{SymInt}(s_1) \;\longrightarrow\; \texttt{SymInt}\!\left(\texttt{sympy.Add}(s_0, s_1)\right) \]

These expression trees flow through every layer of the stack:

• Dynamo tracing — tensor shape attributes become SymInt; operations on shapes produce new SymPy expressions; comparisons produce SymBool (which forces guard generation).
• FakeTensor — the lightweight “meta” tensor used during tracing carries symbolic sizes and strides. Every operator’s shape propagation function (its meta function) computes output shapes as SymPy expressions. For torch.cat([x, y], dim=0), the meta function emits \(s_{x_0} + s_{y_0}\).
• AOTAutograd — preserves symbolic shapes through joint graph tracing; no new symbols are created here.
• TorchInductor — loop bounds and index expressions in Inductor’s loop-level IR are SymPy expressions.

3.2 SymPy Assumptions and Simplification

When a symbol \(s_k\) is created, ShapeEnv configures SymPy assumptions on it:

s_k = sympy.Symbol(f"s{k}", positive=True, integer=True)

These assumptions unlock a range of SymPy simplification rules:

SymPy performs these simplifications eagerly at expression construction time, before any guard is generated. This keeps guard expressions compact — a critical property since guards are compiled to Python source strings that run on every invocation.

3.3 User-Visible Types: SymInt, SymFloat, SymBool

PyTorch exposes three user-visible symbolic types that mirror their Python counterparts:

Definition (Symbolic Types).

• torch.SymInt — wraps a SymNode over a SymPy expression in \(\mathbb{Z}\). Arithmetic operations produce new SymInt objects.
• torch.SymFloat — wraps a SymNode over a SymPy expression in \(\mathbb{R}\). Produced by operations like 1.0 / SymInt.
• torch.SymBool — wraps a SymNode over a SymPy boolean expression. Produced by comparison operators; materializing a SymBool to a concrete bool forces guard generation.

The key invariant: no operation on SymInt/SymFloat/SymBool produces a Python int/float/bool silently. Conversion to a concrete type is always explicit (via int(symint), which records a guard) or via shape_env.evaluate_expr.

4. Guard Generation

📐 Guards are the formal mechanism by which symbolic compilation remains sound. Every decision made during tracing that depends on a symbolic value — taking a branch, selecting a kernel, choosing a layout — must be recorded so the compiled artifact can verify those same decisions are valid at runtime.

4.1 What Triggers a Guard

Guards are generated whenever a SymBool is materialized to a concrete bool. The primary sources are:

Shape attribute access. Accessing x.shape[0] on a tensor with a static (specialized) dimension generates an EQUALS_MATCH guard: L['x'].size()[0] == c. No guard is generated for a dimension that is a free symbol — the symbol itself represents the value.

Explicit comparisons. Dynamo traces Python code symbolically. When it encounters:

if x.shape[0] > 1:
    ...

the condition SymBool(s0 > 1) is evaluated at the hint value (say, 64 → True), the True branch is taken, and the guard s0 > 1 is recorded. On a future call with \(s_0 = 0\), the guard fails → recompile.

Equality assertions. assert x.shape[0] == y.shape[0] during tracing generates the guard s0 == s1. ShapeEnv can also use this to substitute \(s_1 \to s_0\) throughout all subsequent expressions, simplifying the guard set.

torch._check(). The function torch._check(expr) is the programmatic guard insertion API. It marks a boolean expression as unconditionally true for the purposes of compilation, either resolving a guard immediately (if it can be proved from current constraints) or inserting a deferred runtime assert.

4.2 Guard Serialization and Source Annotations

Every sympy.Symbol is associated at creation time with a source — a small object describing how to reach the runtime value that instantiates the symbol. Sources form a compositional hierarchy:

Definition (Guard Serialization). ShapeEnv.produce_guards(placeholders, sources) converts each accumulated ShapeGuard into a Python expression string. For a guard s0 > 1 where \(s_0\) was sourced from GetItemSource(AttrSource(LocalSource("x"), "shape"), 0), the serialized form is:

"L['x'].size()[0] > 1"

This string is compiled into the guard function that runs before each cached graph lookup. The L dictionary is the local variable scope of the traced frame, populated by Dynamo’s frame restoration logic.

produce_guards_verbose additionally emits comment strings with stack traces pointing to where each guard was created — invaluable for debugging spurious recompilations.

4.3 Guard Types and Their Costs

Guards are evaluated in order of decreasing cheapness. EQUALS_MATCH guards appear first; SHAPE_ENV guards (which may involve more complex arithmetic) appear later.

4.4 Divisibility Guards

Division of a symbolic integer requires care. s0 // k is a valid SymPy expression even without a divisibility assumption, but if the compiled code needs the exact floor-division result to control branching or indexing, Dynamo must record that the divisibility holds:

\[ s_0 \,\%\, k = 0 \quad\text{(divisibility guard)} \]

Without this guard, s0 // 2 with \(s_0 = 5\) would produce 2 at trace time (hint-based evaluation) but the guard s0 % 2 == 0 would have been recorded, causing a recompile when \(s_0\) is odd. Dynamo emits these guards automatically when floor-division is applied to a symbol in a context where the result matters (e.g., it is used as a range bound or compared).

5. Specialization on 0 and 1

5.1 Why the Special Cases Exist

Even in dynamic=True mode, Dynamo specializes dimensions that are observed as 0 or 1 at trace time. This is not arbitrary conservatism — both values have deep implications for program semantics.

The case of 0 (zero-size dimensions). A zero-size tensor dimension means numel() == 0. The canonical safe pattern for any tensor operation is:

if numel == 0:
    return torch.empty_like(output)

If Dynamo were to keep \(s_0 = 0\) symbolic, the compiled code would need to handle this branch for all symbol values — generating dead-code-heavy kernels, or requiring a guard s0 != 0 that then specializes the case after all. Dynamo instead records an EQUALS_MATCH guard s == 0 and compiles the zero-size path. Any non-zero input triggers a recompile.

The case of 1 (unit dimensions). A size-1 dimension interacts with broadcasting. Consider x.expand(s0, s1) where \(s_1 = 1\) at trace time. The expand operation is a no-op (a stride-0 view) when \(s_1 = 1\); it becomes a real copy (or a different kernel) when \(s_1 > 1\). Specializing \(s_1 = 1\) allows Inductor to generate the simple no-copy path. Keeping \(s_1\) symbolic would require a branch inside the kernel.

More generally, many PyTorch operators have special-cased fast paths for unit dimensions (e.g., batch normalization with a single-element batch), and these fast paths are not semantically equivalent to the general case in terms of numerical behavior (e.g., batch norm variance is ill-defined for a single element).

Definition (Zero-One Specialization). ShapeEnv’s default policy: if var_to_val[s_k] in {0, 1} at symbol creation time, treat \(s_k\) as a concrete integer. The default range for non-zero-one dynamic dimensions is \([2, \infty)\).

5.2 Performance Implication

The performance implication is easy to miss: switching a batch size from 1 (common in single-sample inference) to 2 triggers a recompile even if dynamic=True is set, because \(B = 1\) was specialized. Users doing single-sample inference with torch.compile should either:

1. Pad the batch to size 2 to avoid the \(B = 1\) specialization, then slice the output; or
2. Compile a separate static artifact for \(B = 1\); or
3. Disable zero-one specialization for the batch dimension (see §5.3).

The same applies to 0-length sequences in masked attention. A sequence of length 0 specializes, so any subsequent call with non-zero length recompiles.

5.3 Disabling Zero-One Specialization

PyTorch provides mark_unbacked (and the associated env var TORCH_COMPILE_UNBACKED_SOURCES) as a path to oblivious-size reasoning — compiling without the 0/1 special cases. However, this requires that the compiled code genuinely handles zero and unit sizes correctly, which many operators do not. The safer approach is to use mark_dynamic with explicit min=2 (or min=1 if the unit case is handled) as described in §7.

6. Backed vs. Unbacked SymInts

6.1 Backed SymInts

A backed symint is a SymInt whose underlying symbol has a known concrete value in ShapeEnv.var_to_val (the hint). Backed symbols arise from tensor shapes: when Dynamo traces x with shape [64, 128] and marks both dimensions dynamic, it allocates \(s_0\) with hint 64 and \(s_1\) with hint 128.

The hint enables:

1. SymPy simplification at trace time. Expressions involving \(s_0\) can be evaluated at the hint to check for zero divisors, select branches, and simplify arithmetic.
2. Static guard generation. The guard function computes the same expression at runtime and checks it matches the expected branch direction.
3. Range initialization. The hint plus mark_dynamic(min=, max=) bounds populate var_to_range.

Backed symbols are named s0, s1, … in diagnostic output.

6.2 Unbacked SymInts

An unbacked symint is a SymInt whose underlying symbol has no concrete hint. These arise from data-dependent operations — computations whose output shape depends on the values of a tensor, not just its shape:

n = x.nonzero().shape[0]          # depends on values of x
k = (x > threshold).sum().item()   # depends on values of x

At trace time, the number of nonzero elements n is unknown because Dynamo uses FakeTensor (a meta/abstract tensor that carries symbolic shape without data). FakeTensor cannot evaluate x.nonzero() for a concrete count.

Definition (Unbacked SymInt). An unbacked symint \(u_k\) is a fresh sympy.Symbol with name u{k} and no entry in var_to_val. Its range is initialized to:

• [0, inf) for size-like contexts (e.g., .nonzero().shape[0])
• (-inf, inf) for general contexts (e.g., .item() on a non-shape value)

Without a hint, SymPy cannot evaluate expressions involving \(u_k\) to a concrete value. This creates two fundamental constraints:

1. No control flow on unbacked symints. if u_k > 5: ... has no hint to guide branch selection. Dynamo cannot take one branch and record the other as a guard — it must raise GuardOnDataDependentSymNode and force a graph break.
2. No static simplification. Expressions like u_k + 1 > u_k cannot be proved True without range information (though range analysis can sometimes resolve this: since \(u_k \geq 0\) and \(u_k + 1 > u_k\) is always true for integers, range analysis may discharge this).

Unbacked symbols are named u0, u1, … in diagnostic output (TORCH_LOGS="+dynamic").

6.3 Deferred Runtime Asserts

Guards on backed symints are checked at the start of each compiled-graph invocation — before any tensor computation runs. But guards involving unbacked symints cannot be checked up front because the unbacked symbol’s value is not computed until the graph runs (it is the output of nonzero, sum, etc.).

Definition (Deferred Runtime Assert). A deferred runtime assert is a SymPy boolean expression involving at least one unbacked symbol, stored in ShapeEnv.deferred_runtime_asserts. The compiler emits these as Python assert statements into the compiled graph body, immediately after the operation that produces the unbacked value:

# Inductor-emitted pseudocode
u0 = aten.nonzero(x).shape[0]
assert u0 >= 0  # deferred runtime assert from ShapeEnv
# ... further computation involving u0 ...

This reinserts the assertion at the point where the value is first available. If the assertion fires at runtime, it indicates that the user’s torch._check calls (see §6.4) made an incorrect claim about the data, and the program is incorrect — not just needing recompilation.

6.4 torch._check and Range Refinement

torch._check(cond) is the mechanism for teaching ShapeEnv about constraints on unbacked symbols. It takes a boolean SymPy expression and:

1. Marks the condition as unconditionally true for all downstream symbolic reasoning in this compilation.
2. Refines the value range of any symbols appearing in cond. For example, torch._check(u0 < 20) tightens \(u_0\)’s range from \([0, \infty)\) to \([0, 19]\).
3. Inserts a deferred runtime assert into the graph, ensuring the claim is validated at execution time.

Example pattern for a data-dependent selection:

indices = (x > threshold).nonzero()
k = indices.shape[0]          # unbacked symint u0
torch._check(k >= 1)          # assert at least one nonzero; tightens range to [1, inf)
torch._check(k <= x.shape[0]) # upper bound; tightens range to [1, s0]
# Now downstream code can use k in bounded range contexts
result = indices[:k, :]       # narrow is safe

def f(x, threshold):
    indices = (x > threshold).nonzero()
    k = indices.shape[0]
    if k > 0:          # <-- problem line
        return x[indices[:, 0]]
    else:
        return x.new_zeros(1)

def f(x, threshold):
    indices = (x > threshold).nonzero()
    k = indices.shape[0]
    # Tell ShapeEnv k >= 0 (already known from range) and assert at runtime
    torch._check(k >= 0)
    # Avoid the branch: use masking or always return a padded tensor
    # Option A: if branch is always safe, use guard_or_false
    # Option B: restructure to avoid control flow on k
    result = x.index_select(0, indices[:, 0]) if k > 0 else x.new_zeros(1)
    # Better: use torch.where or always execute both paths
    out = torch.where(k > 0, x.index_select(0, indices[:, 0]),
                      x.new_zeros(1))  # still hits the same issue
    # Cleanest: use torch.cond (explicit if-else in the graph)
    return torch.cond(k > 0, lambda: x[indices[:, 0]], lambda: x.new_zeros(1))

7. The mark_dynamic / mark_static / torch.export.Dim APIs

7.1 mark_dynamic

torch._dynamo.mark_dynamic(tensor, dim, min=None, max=None) proactively marks a specific tensor dimension as symbolic before compilation. It must be called before the function is first traced by Dynamo.

x = torch.randn(64, 128)
torch._dynamo.mark_dynamic(x, 0)        # s0 = dynamic, s1 = static (=128)
torch._dynamo.mark_dynamic(x, 0, min=1, max=512)  # with explicit range

The effect: when create_symbolic_sizes_strides_storage_offset processes x, dimension 0 is treated as a fresh symbol regardless of its concrete value, even if that value is 0 or 1 (overriding the zero-one specialization policy). The min and max arguments directly initialize var_to_range[s_k] = ValueRanges(min, max) rather than using the default [2, inf).

A subtle point: mark_dynamic bypasses zero-one specialization for the marked dimension. If you mark a batch dimension as dynamic with min=1, then single-sample inference (\(B = 1\)) and multi-sample inference (\(B = 32\)) both use the same compiled artifact — no recompile when batch size changes from 1 to anything else.

7.2 mark_static

torch._dynamo.mark_static(tensor, dim) forces specialization of a dimension even in dynamic=True mode. It exists for dimensions that the user knows are invariant (e.g., the model’s hidden dimension is always 768).

torch._dynamo.mark_static(x, 1)  # always concrete, never symbolic

The effect: in create_symbolic_sizes_strides_storage_offset, dimension 1 is treated as a concrete integer and no symbol is allocated. This restores the full static optimization benefits for that dimension: constant folding, exact tile size selection, loop unrolling.

7.3 torch.export.Dim

torch.export.Dim is the analogous mechanism for torch.export.export. It creates a named symbolic dimension with explicit semantics:

from torch.export import Dim

batch = Dim("batch", min=1, max=512)
seq_len = Dim("seq_len", min=1, max=2048)

ep = torch.export.export(
    model,
    (x, y),
    dynamic_shapes={"x": {0: batch, 1: seq_len}, "y": {0: batch}}
)

Key properties:

• The string name ("batch", "seq_len") is used for readable symbol names in diagnostics and serialized programs.
• Two dimensions annotated with the same Dim object are constrained to be equal — ShapeEnv generates \(s_{\text{batch}_x} = s_{\text{batch}_y}\) and substitutes one for the other.
• Dimensions annotated with different Dim objects are treated as independent symbols, even if they happen to have the same value at trace time.
• Dim.AUTO allocates a symbol with default range \([2, \infty)\) and specializes it if all guards during tracing are consistent with a single concrete value.
• Dim.STATIC forces specialization (analogous to mark_static).
• Dim.DYNAMIC allocates a symbol and raises an error if specialization is attempted (stricter than Dim.AUTO).

Definition (Named Constraint Equations). If dimensions \(d_i\) and \(d_j\) share the same Dim object, ShapeEnv adds the constraint \(s_i = s_j\) to guards, which enables substitution \(s_j \to s_i\) in all subsequent expressions, reducing the symbolic dimension of the problem.

7.4 How Annotations Feed into ShapeEnv

The flow is:

flowchart TD
    A["User annotation
mark_dynamic / mark_static / Dim"] --> B["Dynamo VariableBuilder
processes tensor locals"]
    B --> C["create_symbolic_sizes_strides_storage_offset
in ShapeEnv"]
    C --> D{"Is dim dynamic?"}
    D -->|"yes"| E["Allocate sympy.Symbol s_k
populate var_to_val, var_to_range"]
    D -->|"no"| F["Use concrete integer
no symbol allocated"]
    E --> G["SymInt wrapping SymNode
flows into FakeTensor metadata"]
    F --> G
    G --> H["Operator meta functions
propagate SymPy expressions"]
    H --> I["Guards accumulated
in ShapeEnv.guards"]

The annotation determines the policy at step D. All downstream machinery — SymPy expression construction, guard accumulation, Inductor codegen — operates identically regardless of whether the annotation was mark_dynamic, mark_static, or Dim.

8. Symbolic Shapes in TorchInductor

8.1 Loop IR with Symbolic Bounds

TorchInductor represents computation as a loop-level IR whose fundamental node types are:

• Pointwise(ranges, inner_fn) — an elementwise kernel over a multi-dimensional iteration space.
• Reduction(ranges, reduction_ranges, inner_fn) — a reduction over reduction_ranges producing an output of shape ranges.

The ranges field is a list of SymPy expressions — not necessarily concrete integers. When dimensions are dynamic, these lists contain symbols:

Pointwise(
    ranges=[s0, s1],       # symbolic: e.g. (batch_size, seq_len)
    inner_fn=lambda i, j: x[i, j] + y[i, j]
)

Every index expression inside inner_fn is computed symbolically. For a simple elementwise add over \((s_0, s_1)\) tensors with row-major strides \((s_1, 1)\), the flat index for element \((i, j)\) is:

\[ \text{idx}(i, j) = i \cdot s_1 + j \]

This SymPy expression is emitted into the generated Triton kernel as a runtime arithmetic expression.

8.2 Triton Kernel Parameterization

When Inductor generates a Triton kernel for a dynamic-shape node, the free symbols become kernel launch arguments. The generated Python wrapper looks like:

def kernel_wrapper(x, y, out, s0, s1):
    grid = (triton.cdiv(s0 * s1, BLOCK_SIZE),)
    fused_add_kernel[grid](
        x, y, out,
        s0, s1,              # passed as tl.constexpr-annotated arguments
        BLOCK_SIZE=256,
    )

Inside the Triton kernel, s0 and s1 are available as Python integers at kernel-call time. The Triton compiler and LLVM backend can then perform late constant folding: if s1 = 768 when the kernel is called, LLVM can fold all occurrences of s1 to the constant 768 in the generated PTX, recovering many of the static optimization benefits.

8.3 Constant Folding for Static Dimensions

When a tensor has mixed static and dynamic dimensions — e.g., shape \((s_0, 768)\) — Inductor detects that the second dimension is a concrete integer and folds it into the kernel as a tl.constexpr:

# Mixed: s0 dynamic, dim_1 = 768 static
Pointwise(ranges=[s0, 768], ...)

Inductor emits DIM_1 = 768 as a tl.constexpr argument. This enables:

• Exact tile size selection along the static dimension (tile to 256 knowing \(768 = 3 \times 256\) exactly).
• Loop unrolling along the static dimension.
• Constant stride arithmetic (stride_1 = 1 is known, no multiply needed).

Key result: mixing static and dynamic dimensions is common and well-supported. The transformer’s \((B, L, D)\) pattern — with \(B\) and \(L\) dynamic, \(D\) static — hits this path and benefits from constant-folded \(D\) while flexibly serving variable batch and sequence lengths.

8.4 Guard Tightening in Inductor

Inductor may add guards beyond what Dynamo generated. Two primary sources:

Divisibility guards for tile alignment. If Inductor chooses BLOCK_SIZE = 128 for a dimension of size \(s_0\), it may require \(s_0 \equiv 0 \pmod{128}\) to avoid a boundary check. This guard (s0 % 128 == 0) is added to the Dynamo guard set and checked on each invocation.

Layout guards. Inductor assumes contiguous layouts by default. If the input tensor’s strides are computed symbolically (e.g., after a transpose), Inductor may add guards asserting specific stride values before selecting a kernel variant.

9. Automatic Dynamic Shapes: dynamic=“auto”

Prior to PyTorch 2.1, users had to manually choose between dynamic=False (maximum specialization, maximum recompilation) and dynamic=True (maximum symbolization, conservative optimizations). PyTorch 2.1 introduced dynamic="auto" as the default, implementing a two-phase strategy:

Phase 1 — First compilation (static assumption). The first call to the compiled function is traced with static shapes (all dimensions specialized). This produces a maximally optimized artifact for the example inputs. Guards are recorded normally.

Phase 2 — Recompilation with automatic dynamism. If the input shapes change and a guard fails, Dynamo does not simply recompile with the new static shapes. Instead, it identifies which dimensions changed between the first and second calls and automatically promotes them to dynamic symbols:

\[ \text{dynamic dim} \leftarrow \{d_i : \text{shape}_1[d_i] \neq \text{shape}_2[d_i]\} \]

The second compilation uses mark_dynamic on those dimensions, producing an artifact that works for all future shapes along those axes. Dimensions that did not change remain specialized.

Definition (Auto Dynamic Policy). In dynamic="auto" mode:

1. Call 1: compile with fully static shapes.
2. Call 2 (shape change): recompile with automatically identified dynamic dims.
3. Call 3+: the second artifact handles all conforming shapes; no further recompilation unless a new dimension axis changes.

This policy requires at most 2 compilations per “shape regime” rather than one per shape tuple. The second compiled artifact is slightly less optimized than a fully static one but handles a family of shapes.

Surprisingly, this two-phase approach outperforms naive dynamic=True in many workloads because the first (static) compilation’s guards are tight enough to be cheap, and most models operate in a single shape regime in practice.

10. Debugging Symbolic Shapes

10.1 TORCH_LOGS=“+dynamic”

The most comprehensive debugging tool. Setting TORCH_LOGS="+dynamic" emits verbose output for every symbolic shape decision made during a compilation:

[DEBUG dynamic] create_symbol s0 = 64  source=x.size()[0]  range=[2, inf)
[DEBUG dynamic] create_symbol s1 = 128 source=x.size()[1]  range=[2, inf)
[DEBUG dynamic] specializing s1 == 128 (zero-one spec disabled, but value seen once)
[DEBUG dynamic] guard: s0 > 1  (from: mymodule.py:47: if x.shape[0] > 1:)
[DEBUG dynamic] guard: s0 % 2 == 0  (divisibility: view reshape)

Each log line identifies the symbol, its source, its range, and the code location that triggered the decision. This is the primary tool for diagnosing:

• Unexpected specializations (why is this dimension static when I expected it to be dynamic?)
• Guard proliferation (why are so many guards being generated?)
• Unbacked symint errors (where is the data-dependent shape coming from?)

10.2 torch._dynamo.explain

torch._dynamo.explain(fn)(*args) returns an ExplainOutput object with:

• .graphs — list of captured FX graph fragments
• .guards — list of all guards generated across all graphs
• .break_reasons — locations where graph breaks occurred

For symbolic shape analysis, inspect .guards to see which dimensions generated SHAPE_ENV guards versus EQUALS_MATCH guards. A dimension that should be dynamic but only shows EQUALS_MATCH guards was specialized despite your intent.

10.3 ShapeEnv as a Test Harness

ShapeEnv can be instantiated directly for unit testing symbolic shape logic:

from torch.fx.experimental.symbolic_shapes import ShapeEnv

shape_env = ShapeEnv()
# Manually create symbols and test guard generation
# Useful for verifying that a specific expression generates the expected guard
guards_str = shape_env.produce_guards(placeholders, sources)
valid = shape_env.evaluate_guards_for_args(placeholders, args)

This is the approach used in PyTorch’s own test suite for symbolic shapes (test/test_symbolic_shapes.py), and is valuable for library authors who want to ensure their operator meta functions produce correct symbolic shapes.

10.4 Common Failure Modes

References