Indexing

Indexing a wrapper follows the same principle as Region views: you index through a region view, so a slice or integer index addresses one region’s axes and the other regions pass through untouched. The index syntax itself is ordinary NumPy/torch indexing — slices, integers, None, Ellipsis; the authoritative reference for that syntax is the NumPy indexing guide. What NEML2 adds is the region scoping and the metadata bookkeeping that comes with it.

Region-scoped indexing

A key on t.dynamic_batch or t.sub_batch addresses only that region’s axes; everything before and after passes through as :. The wrapper’s sub_batch_ndim is adjusted automatically — an integer index drops an axis, None inserts one:

import torch
from neml2.types import Scalar

V = Scalar(torch.arange(5.0)).sub_batch.retag(1)  # length-5 sub-batch axis
print(
    "V.sub_batch[0]   ->",
    tuple(V.sub_batch[0].data.shape),
    "sub_ndim",
    V.sub_batch[0].sub_batch_ndim,
)
print("V.sub_batch[1:]  ->", tuple(V.sub_batch[1:].data.shape))
print(
    "V.sub_batch[None]->",
    tuple(V.sub_batch[None].data.shape),
    "sub_ndim",
    V.sub_batch[None].sub_batch_ndim,
)

V.sub_batch[0]   -> () sub_ndim 0
V.sub_batch[1:]  -> (4,)
V.sub_batch[None]-> (1, 5) sub_ndim 2

The dynamic-batch region indexes the same way and preserves the sub-batch count:

from neml2.types import Vec

field = Vec(torch.randn(4, 3))  # 4 dynamic states, base (3,)
print("dynamic_batch[1:] ->", tuple(field.dynamic_batch[1:].data.shape))
print("dynamic_batch[0]  ->", tuple(field.dynamic_batch[0].data.shape))
print("batch[::2]        ->", tuple(field.batch[::2].data.shape))

dynamic_batch[1:] -> (3, 3)
dynamic_batch[0]  -> (3,)
batch[::2]        -> (2, 3)

Ellipsis and None

Inside a region, ... expands to as many : as the region has axes (at most one per key), and None inserts a fresh axis — the same meanings the NumPy guide gives them, just confined to the region:

grid = Scalar(torch.randn(3, 4)).sub_batch.retag(2)  # two sub-batch axes
print(
    "sub_batch[..., 0] ->",
    tuple(grid.sub_batch[..., 0].data.shape),
    "sub_ndim",
    grid.sub_batch[..., 0].sub_batch_ndim,
)
print(
    "sub_batch[None]   ->",
    tuple(grid.sub_batch[None].data.shape),
    "sub_ndim",
    grid.sub_batch[None].sub_batch_ndim,
)

sub_batch[..., 0] -> (3,) sub_ndim 1
sub_batch[None]   -> (1, 3, 4) sub_ndim 3

Base components

The base region of a fixed-base primitive is not indexable through a view — its shape and packing are part of the type, so reaching into it by raw position would bypass the convention (the Mandel √2 on an SR2, for instance). To pull a component out as a Scalar, use the vec_component free function:

from neml2.types import vec_component

v = Vec.fill(6.0, 4.0, 0.0)
x = vec_component(v, 0)
print("vec_component(v, 0) ->", type(x).__name__, x.data)

vec_component(v, 0) -> Scalar tensor(6., dtype=torch.float64)

The dynamic-base Tensor is the exception: because its base rank is a runtime field rather than a fixed convention, its base view is indexable, which the equation-system / solver layer relies on to slice Jacobian blocks:

from neml2.types import Tensor

blk = Tensor(torch.randn(2, 3, 4), batch_ndim=1)  # base (3, 4)
print("base[..., -1] ->", tuple(blk.base[..., -1].data.shape))
print("base[1:3, :]  ->", tuple(blk.base[1:3, :].data.shape))

base[..., -1] -> (2, 3)
base[1:3, :]  -> (2, 2, 4)

How this differs from plain torch indexing

On a raw torch.Tensor you index the flat shape and carry the responsibility of knowing which axis is which. Through a region view you index one region, the rest pass through automatically, and the returned wrapper’s region metadata is updated for you — so a sub-batch slice stays a valid sub-batch, and a dropped axis decrements sub_batch_ndim instead of silently misaligning later operations. For the index syntax (basic slicing, integer arrays, None, Ellipsis), the NumPy indexing guide is the reference.

See also