Tensor shape regions

Open in Colab

A raw torch.Tensor has one flat shape. A NEML2 wrapper gives that shape structure: the trailing axes are the fixed mathematical base, and the leading axes are split into two batch regions that the framework treats differently. Understanding the split is the key to reading any wrapper’s shape.

The decomposition

A wrapper’s data.shape reads left-to-right as three contiguous regions:

data.shape == (*dynamic_batch_shape, *sub_batch_shape, *base_shape)
              └──── leading ────┘└─ middle (static) ─┘└─ BASE_NDIM ─┘

  • Base shape — the trailing BASE_NDIM axes, fixed by the wrapper type ((6,) for SR2, () for Scalar). These encode the mathematical structure and never broadcast.

  • Sub-batch shape — a small, static batching region for per-site structure (lookup-table axes, finite-volume cells, slip systems). The chain-rule machinery treats these axes specially. Default 0 — most models don’t need it.

  • Dynamic batch shape — everything left over. Free-form, sized at call time, and traced as dynamic by torch.export so one compiled artifact handles every batch size.

Reading the regions off a wrapper

Each region has an accessor. With no sub-batch axes declared, everything that isn’t base is dynamic batch:

import torch
from neml2.types import SR2, Scalar

grid = SR2(torch.randn(50, 200, 6))
print("data.shape          ", tuple(grid.data.shape))
print("batch_shape         ", tuple(grid.batch_shape))
print("dynamic_batch_shape ", tuple(grid.dynamic_batch_shape))
print("sub_batch_shape     ", tuple(grid.sub_batch_shape))
print("base_shape          ", tuple(grid.base_shape))

data.shape           (50, 200, 6)
batch_shape          (50, 200)
dynamic_batch_shape  (50, 200)
sub_batch_shape      ()
base_shape           (6,)

batch_shape is the combined dynamic_batch + sub_batch region — “everything that isn’t base”. To create a sub-batch region, mark how many trailing batch axes are per-site with .sub_batch.retag(n):

field = Scalar(torch.randn(4, 20))  # 4 dynamic states × a length-20 axis
field = field.sub_batch.retag(1)  # mark the trailing axis as sub-batch
print("batch_shape         ", tuple(field.batch_shape))
print("dynamic_batch_shape ", tuple(field.dynamic_batch_shape))
print("sub_batch_shape     ", tuple(field.sub_batch_shape))
print("sub_batch_ndim      ", field.sub_batch_ndim)

batch_shape          (4, 20)
dynamic_batch_shape  (4,)
sub_batch_shape      (20,)
sub_batch_ndim       1

Dynamic vs sub-batch

The two batch regions exist precisely because the framework treats them differently:

Region

Sized at…

Traced as…

Broadcasts with…

Typical use

Dynamic batch

call time

dynamic dim

everything

every ordinary batch — N material points, time steps

Sub-batch

construction time

static shape

other sub-batches of matching width

per-site structure — interpolation-table axis, FV cell, slip-system axis

Default sub_batch_ndim = 0: a model that doesn’t need the distinction ignores it, and everything sits in the dynamic region. Sub-batch axes do not participate in dynamic-batch broadcasting; they behave like a small extra structural region the chain rule accumulates over. The batching rules are covered in Batching and broadcasting.

Two regions you usually don’t construct

An internal K region. Chain-rule tangents carry an extra leading K (seed-direction) region to the left of the dynamic batch. Primal values — everything you build by hand — have k_ndim == 0, so the region is absent and you can ignore it; it is managed entirely by the framework’s derivative machinery.

print("k_ndim of a hand-built wrapper:", grid.k_ndim)

k_ndim of a hand-built wrapper: 0

A dynamic-base Tensor. Alongside the fixed-base primitives there is a Tensor type whose base rank is a runtime field, not a class invariant. It is the storage unit for variable-shape Jacobian blocks at the equation-system / solver layer; the same class can stand in for a scalar, an SR2-shaped block, or an arbitrary (rows, cols) matrix. You rarely construct one directly, but it carries the same region vocabulary with batch_ndim / sub_batch_ndim as explicit fields:

from neml2.types import Tensor

block = Tensor.zeros(batch_shape=(8,), base_shape=(6, 6))
print("data.shape ", tuple(block.data.shape))
print("batch_ndim ", block.batch_ndim, " base_ndim ", block.base_ndim)

data.shape  (8, 6, 6)
batch_ndim  1  base_ndim  2

See also