Tensor types

Currently, PyTorch is the only supported tensor backend in NEML2. Therefore, all tensor types in NEML2 directly inherit from torch::Tensor. In the future, support for other tensor backends may be added, but the public-facing interfaces will remain largely the same.

Dynamically shaped tensor

BatchTensor is a general-purpose dynamically shaped tensor type for batched tensors. With a view towards vectorization, the same set of operations can be "simultaneously" applied to a "batch" of (logically the same) tensors. To provide a unified user interface for dealing with such batched operation, NEML2 assumes that the first \(N\) dimensions of a tensor are batched dimensions, and the following dimensions are the base (logical) dimensions.

‍Unlike PyTorch, NEML2 explicitly distinguishes between batch dimensions and base (logical) dimensions.

A BatchTensor can be created using

BatchTensor A(torch::rand({1, 1, 5, 2}), 2);

where A is a tensor with 2 batch dimensions. The batch sizes of A is (1, 1):

auto batch_sz = A.batch_sizes();

neml2_assert(batch_sz == {1, 1});

and the base (logical) sizes of A is (5, 2):

auto base_sz = A.base_sizes();

neml2_assert(batch_sz == {5, 2});

Statically shaped tensors

FixedDimTensor is the parent class for all the tensor types with a fixed base shape. It is templated on the base shape of the tensor. NEML2 offers a rich collection of primitive tensor types inherited from FixedDimTensor. Currently implemented primitive tensor types are summarized below

Tensor type	Base shape	Description
Scalar	\(()\)	Rank-0 tensor, i.e. scalar
Vec	\((3)\)	Rank-1 tensor, i.e. vector
R2	\((3,3)\)	Rank-2 tensor
SR2	\((6)\)	Symmetric rank-2 tensor
WR2	\((3)\)	Skew-symmetric rank-2 tensor
R3	\((3,3,3)\)	Rank-3 tensor
SFR3	\((6,3)\)	Rank-3 tensor with symmetry on base dimensions 0 and 1
R4	\((3,3,3,3)\)	Rank-4 tensor
SSR4	\((6,6)\)	Rank-4 tensor with minor symmetry
R5	\((3,3,3,3,3)\)	Rank-5 tensor
SSFR5	\((6,6,3)\)	Rank-5 tensor with minor symmetry on base dimensions 0-3
Rot	\((3)\)	Rotation tensor represented in the Rodrigues form
Quaternion	\((4)\)	Quaternion
MillerIndex	\((3)\)	Crystal direction or lattice plane represented as Miller indices

Furthermore, all primitive tensor types can be "registered" as variables on a LabeledAxis, which will be discussed in the following section on labeled view.

Working with tensors

Tensor creation

A factory tensor creation function produces a new tensor. All factory functions adhere to the same schema:

<TensorType>::<function_name>(<function-specific-options>, const torch::TensorOptions & options);

where <TensorType> is the class name of the primitive tensor type listed above, and <function-name> is the name of the factory function which produces the new tensor. <function-specific-options> are any required or optional arguments a particular factory function accepts. Refer to each tensor type's class documentation for the concrete signature. The last argument const torch::TensorOptions & options configures the data type, device, layout and other "meta" properties of the produced tensor. The commonly used meta properties are

dtype: the data type of the elements stored in the tensor. Available options are kUInt8, kInt8, kInt16, kInt32, kInt64, kFloat32, and kFloat64.
layout: the striding of the tensor. Available options are kStrided (dense) and kSparse.
device: the compute device where the tensor will be allocated. Available options are kCPU and kCUDA.
requires_grad: whether the tensor is part of a function graph used by automatic differentiation to track functional relationship. Available options are true and false.

For example, the following code

auto a = SR2::zeros({5, 3},
                    torch::TensorOptions()
                      .device(torch::kCPU)
                      .layout(torch::kStrided)
                      .dtype(torch::kFloat32));

creates a statically (base) shaped, dense, single precision tensor of type SR2 filled with zeros, with batch shape \((5, 3)\), allocated on the CPU.

Tensor broadcasting

Quoting Numpy's definition of broadcasting:

‍The term broadcasting describes how NumPy treats arrays with different shapes during arithmetic operations. Subject to certain constraints, the smaller array is “broadcast” across the larger array so that they have compatible shapes.

NEML2's broadcasting semantics is largely the same as those of Numpy and PyTorch. However, since NEML2 explicitly distinguishes between batch and base dimensions, the broadcasting semantics must also be extended. Two NEML2 tensors are said to be batch-broadcastable if iterating backward from the last batch dimension, one of the following is satisfied:

Both tensors have the same size on the dimension;
One tensor has size 1 on the dimension;
The dimension does not exist in one tensor.

Base-broadcastable follows a similar definition. Most binary operators on dynamically shaped tensors, i.e., those of type BatchTensor, require the operands to be both batch- and base-broadcastable. On the other hand, most binary operators on statically base shaped tensors, i.e., those of pritimitive tensor types, only require the operands to be batch-broadcastable.

Tensor indexing

In defining the forward operator of a material model, many logically different tensors representing inputs, outputs, residuals, and Jacobians have to be created, copied, and destroyed on the fly. These operations occupy a significant amount of computing time, especially on GPUs.

To address this challenge, NEML2 creates views, instead of copies, of tensors whenever possible. As its name suggests, the view of a tensor is a possibly different interpretation of the underlying data. Quoting the PyTorch documentation:

‍For a tensor to be viewed, the new view size must be compatible with its original size and stride, i.e., each new view dimension must either be a subspace of an original dimension, or only span across original dimensions \(d, d+1, ..., d+k\) that satisfy the following contiguity-like condition that \(\forall i = d,...,d+k-1\),
\[ \text{stride}[i] = \text{stride}[i+1] \times \text{size}[i+1] \]

Otherwise, it will not be possible to view self tensor as shape without copying it.

In NEML2, use base_index for indexing the base dimensions and batch_index for indexing the batch dimensions:

using namespace torch::indexing;
BatchTensor A(torch::tensor({{2, 3, 4}, {-1, -2, 3}, {6, 9, 7}}), 1);
// A = [[  2  3  4]
//      [ -1 -2  3]
//      [  6  9  7]]
BatchTensor B = A.batch_index({Slice(0, 2)});
// B = [[  2  3  4]
//      [ -1 -2  3]]
BatchTensor C = A.base_index({Slice(1, 3)});
// C = [[  3  4]
//      [ -2  3]
//      [  9  7]]

To modify the content of a tensor, use base_index_put or batch_index_put:

A.base_index_put({Slice(1, 3)}, torch::ones({3, 2}));
// A = [[  2  1  1]
//      [ -1  1  1]
//      [  6  1  1]]
A.batch_index_put({Slice(0, 2)}, torch::zeros({2, 3}));
// A = [[  0  0  0]
//      [  0  0  0]
//      [  6  1  1]]

A detailed explanation on tensor indexing APIs is available as part of the official PyTorch documentation.

Tensor labeling

In the context of material modeling, oftentimes views of tensors have practical/physical meanings. For example, given a logically 1D tensor with base size 9, its underlying data in an arbitrary batch may look like

equivalent plastic strain   2.1
            cauchy stress  -2.1
                              0
                            1.3
                           -1.1
                            2.5
                            2.5
              temperature 102.9
                     time   3.6

where component 0 stores the scalar-valued equivalent plastic strain, components 1-6 store the tensor-valued cauchy stress (we use the Mandel notation for symmetric second order tensors), component 7 stores the scalar-valued temperature, and component 8 stores the scalar-valued time.

The string indicating the physical meaning of the view, e.g., "cauchy stress", is called a "label", and the view of the tensor indexed by a label is called a "labeled view", i.e.,

cauchy stress  -2.1
                  0
                1.3
               -1.1
                2.5
                2.5

NEML2 provides a data structure named LabeledAxis to facilitate the creation and modification of labels, and a data structure named LabeledTensor to facilitate the creation and modification of labeled views.

The LabeledAxis contains all information regarding how an axis of a LabeledTensor is labeled. The following naming convention is used:

Item: A labelable slice of data
Variable: An item that is also of a NEML2 primitive tensor type
Sub-axis: An item of type LabeledAxis

So yes, an axis can be labeled recursively, e.g.,

     0 1 2 3 4 5     6     7 8 9 10 11 12   13   14
/// |-----------| |-----| |              | |  | |  |
///       a          b    |              | |  | |  |
/// |-------------------| |--------------| |--| |--|
///          sub                  a          b    c

The above example represents an axis of size 15. This axis has 4 items: a, b, c, and sub.

"a" is a variable of storage size 6 (possibly of type SR2).
"b" is a variable of type Scalar.
"c" is a variable of type Scalar.
"sub" is a sub-axis of type LabeledAxis. "sub" by itself represents an axis of size 7, containing 2 items:
- "a" is a variable of storage size 6.
- "b" is a variable of type Scalar.

Duplicate labels are not allowed on the same level of the axis, e.g. "a", "b", "c", and "sub" share the same level and so must be different. However, items on different levels of an axis can share the same label, e.g., "a" on the sub-axis "sub" has the same label as "a" on the main axis. In NEML2 convention, item names are always fully qualified, and a sub-axis is prefixed with a left slash, e.g. item "b" on the sub-axis "sub" can be denoted as "sub/b" on the main axis.

‍A label cannot contain: white spaces, quotes, left slash (/), or new line.

Due to performance considerations, a LabeledAxis can only be modified, e.g., adding/removing variables and sub-axis, at the time a model is constructed. After the model construction phase, the LabeledAxis associated with that model can no longer be modified over the entire course of the simulation.

Refer to the documentation for a complete list of APIs for creating and modifying a LabeledAxis.

LabeledTensor is the primary data structure in NEML2 for working with labeled tensor views. Each LabeledTensor consists of one BatchTensor and one or more LabeledAxiss. The LabeledTensor is templated on the base dimension \(D\). LabeledVector and LabeledMatrix are the two most widely used data structures in NEML2.

LabeledTensor handles the creation, modification, and accessing of labeled tensors. Recall that all primitive data types in a labeled tensor are flattened, e.g., a symmetric fourth order tensor of type SSR4 with batch size (5) and base size (6, 6) are flattened to have base size (36) in the labeled tensor. The documentation provides a complete list of APIs. The commonly used methods are

operator() for retrieving a labeled view into the raw (flattened) data without reshaping
get for retrieving a labeled view and reshaping it to the correct shape
set for setting values for a labeled view
slice for slicing a sub-axis along a specific base dimension
block for sub-indexing the LabeledTensor with \(D\) sub-axis names

Table of Contents