Batching and broadcasting¶

Vectorizing a model over many material points is the whole point of the tensor backend: the same forward body evaluates one state or a million, and the only difference is the shape of the input. This page shows the batching pattern, the broadcasting rules wrappers follow, and the one batching concept torch has no equivalent for — the sub-batch region.

One body, any batch shape¶

A model’s forward dispatches one PyTorch kernel per operation, and those kernels broadcast across leading axes natively. So we can hand the same model a single strain, a 1-D batch, or a 2-D grid and only the leading axes change. First, a small elasticity model to evaluate:

%%writefile elasticity.i
[Models]
  [elasticity]
    type = LinearIsotropicElasticity
    coefficients      = '200e3          0.3'
    coefficient_types = 'YOUNGS_MODULUS POISSONS_RATIO'
  []
[]

Writing elasticity.i

import torch
import neml2
from neml2.types import SR2

model = neml2.load_model("elasticity.i", "elasticity")

single = SR2.fill(0.01, 0, 0, 0, 0, 0)  # base only, no batch axes
batch = SR2(torch.randn(10_000, 6) * 0.01)  # one leading batch dim
grid = SR2(torch.randn(50, 200, 6) * 0.01)  # two leading batch dims

print("single ->", tuple(model(single).data.shape))
print("batch  ->", tuple(model(batch).data.shape))
print("grid   ->", tuple(model(grid).data.shape))

single -> (6,)
batch  -> (10000, 6)
grid   -> (50, 200, 6)

Leading batch dims are completely free-form. A Python loop around a single-state call leaves a lot of throughput on the table — pass the whole batch as one tensor whenever you can. Vectorization shows the timing difference.

Broadcasting¶

Binary operators between wrappers broadcast their batch regions using PyTorch’s standard rules: right-aligned by axis, size-1 axes stretched to match, mismatched non-1 sizes an error. A Scalar broadcasts against any wrapper’s base, and a base-only value broadcasts against a batched one:

from neml2.types import Scalar

pressure = Scalar(torch.linspace(1.0, 5.0, 5))  # shape (5,)
unit = SR2.fill(1, 1, 1, 0, 0, 0)  # base only, shape (6,)
scaled = pressure * unit  # (5,) Scalar × (6,) SR2
print("Scalar(5) * SR2(base) ->", tuple(scaled.data.shape))

offset = SR2(torch.randn(5, 6)) + unit  # (5, 6) + (6,)
print("SR2(5) + SR2(base)    ->", tuple(offset.data.shape))

Scalar(5) * SR2(base) -> (5, 6)
SR2(5) + SR2(base)    -> (5, 6)

Sub-batch metadata survives the op¶

Algebraic operators don’t just broadcast the raw tensors — they unify the operands’ sub-batch widths and tag the result accordingly. So a “global” parameter and a “per-site” field combine cleanly at any dynamic batch size, and the result still knows which axis is per-site:

glob = Scalar(2.0)  # a global scalar, no sub-batch
per_site = Scalar(torch.randn(20)).sub_batch.retag(1)  # length-20 sub-batch
combined = glob * per_site
print("sub_batch_shape ->", tuple(combined.sub_batch_shape))
print("sub_batch_ndim  ->", combined.sub_batch_ndim)

sub_batch_shape -> (20,)
sub_batch_ndim  -> 1

Declaring sub-batch axes in an input file¶

.sub_batch.retag(n) is also valid inside a [Tensors] HIT block, whose expr is evaluated as Python. A lookup-table control axis is the canonical case — here 20 interpolation points marked as the sub-batch:

%%writefile controls.i
[Tensors]
  [T_controls]
    type = Python
    expr = 'linspace(Scalar(300.0).sub_batch, Scalar(1200.0).sub_batch, 20)'
  []
[]

Writing controls.i

controls = neml2.load_input("controls.i").get_tensor("T_controls")
print("sub_batch_shape ->", tuple(controls.sub_batch_shape))
print("sub_batch_ndim  ->", controls.sub_batch_ndim)

sub_batch_shape -> (20,)
sub_batch_ndim  -> 1

Relation to torch¶

The leading-batch broadcasting above is PyTorch’s — wrappers add no new rule there; they just keep the base packing intact through it. The sub-batch region is the genuinely new concept: a torch tensor has no way to say “this axis is per-site structure, not an ordinary batch”, which is exactly the distinction the chain rule needs when models on different sites compose.