Region views

On a raw torch.Tensor, t.unsqueeze(0) or torch.sum(t, dim=1) address axes by absolute position — and once a tensor mixes batch, sub-batch, and base axes, an absolute index is ambiguous and easy to get wrong. NEML2 fixes this by exposing every shape-manipulating op through a region view, so each reshape, reduction, or concatenation names the region it acts on.

The views

Four properties select a region; each returns a fresh wrapper, so calls chain:

  • t.dynamic_batch — the dynamic batch region. Ops preserve sub_batch_ndim.

  • t.sub_batch — the sub-batch region. Ops adjust sub_batch_ndim.

  • t.batch — the combined region, for read-only .shape / .ndim and for cat; shape-changing ops on it raise (see below).

  • t.base — the base region. Read-only except transpose on square-base types (R2, SSR4).

Every mutable view exposes .shape, .ndim, .unsqueeze(dim), .squeeze(dim), and .expand(*shape):

from neml2.types import SR2, R2

broadcast = SR2.fill(0.1, -0.05, -0.05, 0, 0, 0).dynamic_batch.expand(20)
print("dynamic_batch.expand(20) ->", tuple(broadcast.data.shape))

tr = R2.identity().base.transpose(-2, -1)
print("base.transpose(-2, -1)   ->", tuple(tr.data.shape))

dynamic_batch.expand(20) -> (20, 6)
base.transpose(-2, -1)   -> (3, 3)

Why t.batch can’t be reshaped

t.batch spans both batch regions at once, so a shape-changing op on it would be ambiguous: should a new axis land in the dynamic region or the sub-batch region? Rather than guess, those ops raise and tell you which view to pick:

try:
    SR2.zeros(3).batch.unsqueeze(-1)
except TypeError as err:
    print(err)

batch.unsqueeze is ambiguous; use t.dynamic_batch.unsqueeze(...) or t.sub_batch.unsqueeze(...)

Sub-batch-only operations

The sub-batch view carries a few extra ops for building and collapsing per-site structure — retag(n) / flatten() to change how many trailing axes count as sub-batch, expand_at(size) to broadcast in a new site axis, and diagonalize() to turn a per-site vector into a per-site diagonal:

import torch
from neml2.types import Scalar

s = Scalar(torch.randn(5)).sub_batch.retag(1)
print("expand_at(3) ->", tuple(s.sub_batch.expand_at(3).data.shape))
print("diagonalize  ->", tuple(s.sub_batch.diagonalize().data.shape))
print("flatten ndim ->", s.sub_batch.flatten().sub_batch_ndim)

expand_at(3) -> (3, 5)
diagonalize  -> (5, 5)
flatten ndim -> 0

Disambiguating the free functions

This is where region views earn their keep. Reductions, concatenation, and interpolation live as free functions in neml2.types, and they take a region view instead of a dim= integer. The view tells the function which region the axis index addresses — and lets it update the wrapper’s metadata correctly (e.g. dropping sub_batch_ndim when a sub-batch axis is reduced away).

sum, mean, and diff reduce or difference one region’s axis:

from neml2.types import sum, mean, diff, linspace, cat, stack

ramp = linspace(Scalar(300.0).sub_batch, Scalar(1200.0).sub_batch, 20)
print("ramp sub_batch_shape ->", tuple(ramp.sub_batch_shape))

total = sum(ramp.sub_batch, 0)  # reduce the sub-batch axis
print("sum over sub_batch    ->", tuple(total.data.shape), "sub_ndim", total.sub_batch_ndim)

avg = mean(ramp.sub_batch, 0)
print("mean over sub_batch   ->", tuple(avg.data.shape))

delta = diff(ramp.sub_batch, n=1, dim=-1)
print("diff along sub_batch  ->", tuple(delta.sub_batch_shape))

ramp sub_batch_shape -> (20,)
sum over sub_batch    -> () sub_ndim 0
mean over sub_batch   -> ()
diff along sub_batch  -> (19,)

stack and cat combine along a region’s axis; the view fixes which region grows:

from neml2.types import Vec

v0 = Vec.fill(6.0, 4.0, 0.0)
v1 = Vec.fill(8.0, 5.0, 0.0)
print("stack on dynamic_batch ->", tuple(stack([v0.dynamic_batch, v1.dynamic_batch]).data.shape))

a = Vec(torch.randn(2, 3))
b = Vec(torch.randn(3, 3))
print("cat on dynamic_batch   ->", tuple(cat([a.dynamic_batch, b.dynamic_batch], dim=0).data.shape))

stack on dynamic_batch -> (2, 3)
cat on dynamic_batch   -> (5, 3)

And linspace / logspace build a ramp along whichever region you pass the endpoints as — the single v3 spelling that replaces v2’s separate dynamic_linspace / intmd_linspace / base_linspace:

dyn = linspace(Scalar(0.0).dynamic_batch, Scalar(1.0).dynamic_batch, 5)
print("endpoints as .dynamic_batch -> new dynamic axis ", tuple(dyn.dynamic_batch_shape))
sub = linspace(Scalar(0.0).sub_batch, Scalar(1.0).sub_batch, 5)
print("endpoints as .sub_batch     -> new sub-batch axis", tuple(sub.sub_batch_shape))

endpoints as .dynamic_batch -> new dynamic axis  (5,)
endpoints as .sub_batch     -> new sub-batch axis (5,)

See also