Primitive (fixed-base-shape) tensor types

Open in Colab

These are the wrappers you construct directly: a Scalar, a 3-vector, a symmetric second-order tensor, and so on. Each one fixes the shape and the packing convention of its trailing (base) axes — that is what separates an SR2 from a bare torch.Tensor of shape (6,).

The class hierarchy

Every primitive shares one base shape per type and inherits its arithmetic and factories from a common intermediate class:

TensorWrapper            (abstract — shape decomposition + region views)
    └── PrimitiveTensor  (concrete intermediate — generic ops + factories)
            ├── Scalar
            └── Vec, R2, SR2, WR2, MRP, Quaternion, MillerIndex, SSR4

PrimitiveTensor is where the generic operators (+, -, *, /, -x) and the shape factories (zeros, ones, full, empty, fill) live. Each concrete leaf adds only its type-specific factories — e.g. R2.identity, SSR4.identity_sym, or SR2.fill’s Mandel-aware overloads.

The catalog

Type

Base shape

Storage / convention

Scalar

()

A single number per batch entry. The wrapper exists so mixed ops like Scalar * SR2 reliably return an SR2.

Vec

(3,)

3-vector.

MRP

(3,)

Modified Rodrigues parameters for a 3D rotation: n · tan(θ/4), zero vector = identity.

MillerIndex

(3,)

Integer crystallographic direction / plane normal, stored as float for differentiability.

Quaternion

(4,)

Unit quaternion orientation, packed (w, x, y, z).

R2

(3, 3)

Full second-order tensor (no symmetry).

WR2

(3,)

Skew-symmetric second-order tensor stored as an axial 3-vector (w0, w1, w2); the 3×3 skew form is recovered by r2_from_wr2.

SR2

(6,)

Symmetric second-order tensor in Mandel notation: (σ₁₁, σ₂₂, σ₃₃, √2·σ₂₃, √2·σ₁₃, √2·σ₁₂). The off-diagonal √2 makes the inner product A : B = a · b in 6-vector form.

SSR4

(6, 6)

Fourth-order tensor with both minor symmetries (the elasticity-tensor symmetry class) in Mandel packing.

import torch
from neml2.types import Scalar, Vec, MRP, MillerIndex, Quaternion, R2, WR2, SR2, SSR4

for cls in (Scalar, Vec, MRP, MillerIndex, Quaternion, R2, WR2, SR2, SSR4):
    print(f"{cls.__name__:12s} BASE_SHAPE = {cls.BASE_SHAPE}")
Scalar       BASE_SHAPE = ()
Vec          BASE_SHAPE = (3,)
MRP          BASE_SHAPE = (3,)
MillerIndex  BASE_SHAPE = (3,)
Quaternion   BASE_SHAPE = (4,)
R2           BASE_SHAPE = (3, 3)
WR2          BASE_SHAPE = (3,)
SR2          BASE_SHAPE = (6,)
SSR4         BASE_SHAPE = (6, 6)

Packing conventions are part of the type

Because the convention lives in the type, the constructor does the packing for you. SR2.fill accepts the Mandel √2 scaling internally, so you pass physical components and get correct storage — the 1-, 3-, and 6-component overloads cover a hydrostatic state, a diagonal state, and a general symmetric state respectively:

print("SR2.fill(p)        ->", tuple(SR2.fill(2.0).data.shape))
print("SR2.fill(d1,d2,d3) ->", tuple(SR2.fill(1.0, 2.0, 3.0).data.shape))
print("SR2.fill(6 comps)  ->", tuple(SR2.fill(1, 2, 3, 4, 5, 6).data.shape))
print("SR2.identity()     ->", SR2.identity().data)
SR2.fill(p)        -> (6,)
SR2.fill(d1,d2,d3) -> (6,)
SR2.fill(6 comps)  -> (6,)
SR2.identity()     -> tensor([1., 1., 1., 0., 0., 0.])

dtype

Scalar defaults to torch.float64 (the precision NEML2’s typed algebra runs in) for both literal construction and its factories. The other wrappers take an optional dtype= and otherwise fall through to torch’s global default:

print("Scalar(200e3).dtype  ->", Scalar(200e3).dtype)
print("Vec.zeros(5).dtype   ->", Vec.zeros(5).dtype)
print("Vec.zeros(5, dtype=float64) ->", Vec.zeros(5, dtype=torch.float64).dtype)
Scalar(200e3).dtype  -> torch.float64
Vec.zeros(5).dtype   -> torch.float32
Vec.zeros(5, dtype=float64) -> torch.float64

Constructing wrappers

Every primitive inherits the same shape-factory family from PrimitiveTensor. Each takes the dynamic batch shape as leading positional args and pads the call to torch with the type’s base shape:

print("Vec.zeros(4)            ->", tuple(Vec.zeros(4).data.shape))
print("SR2.full(2, 3, fill=0.1)->", tuple(SR2.full(2, 3, fill_value=0.1).data.shape))
print("R2.identity()          ->", tuple(R2.identity().data.shape))
print("SSR4.identity_sym()    ->", tuple(SSR4.identity_sym().data.shape))
Vec.zeros(4)            -> (4, 3)
SR2.full(2, 3, fill=0.1)-> (2, 3, 6)
R2.identity()          -> (3, 3)
SSR4.identity_sym()    -> (6, 6)

Scalar adds the torch-analogue creation helpers, keeping the float64 default:

  • Scalar(<float>) / Scalar(<list>) — direct literal coercion.

  • Scalar.zeros / Scalar.ones / Scalar.fullfloat64 shape factories.

  • Scalar.arange(start, end, step) — mirrors torch.arange.

  • Scalar.from_value(x, like=other) — promote a Python number, inheriting dtype / device from an existing wrapper (handy inside a leaf’s forward to build an in-place neutral).

To build a ramp between two endpoints, use the Region views linspace / logspace free functions rather than a Scalar method — the target region is chosen by which region view you pass the endpoints as.

print("Scalar([1, 2, 3]) ->", Scalar([1, 2, 3]).data)
print("Scalar.arange(0, 1, 0.25) ->", Scalar.arange(0.0, 1.0, 0.25).data)
ref = SR2.zeros(3, dtype=torch.float64)
print("from_value inherits dtype ->", Scalar.from_value(0.5, like=ref).dtype)
Scalar([1, 2, 3]) -> tensor([1., 2., 3.], dtype=torch.float64)
Scalar.arange(0, 1, 0.25) -> tensor([0.0000, 0.2500, 0.5000, 0.7500], dtype=torch.float64)
from_value inherits dtype -> torch.float64

See also

  • Tensor shape regions — how the leading axes split into batch regions once you start batching these types.

  • Syntax catalog — the per-type option reference for every registered object.