NEML2 tensor indexing works much in the same way as in NumPy and PyTorch. We also offer a simple one-to-one translation between C++ and Python tensor indexing API. The major difference between indexing (batched) NEML2 tensor and indexing tensors in PyTorch or N-D arrays in NumPy is that most indexing APIs have two variants – one for indexing batch dimensions and the other for indexing base dimensions.
Single element indexing
Single element indexing works exactly like that for indexing flat containers such as vectors, arrays, lists, etc. The indexing is 0-based and accepts negative integers (for reverse indexing).
Single element indexing can be used to index multidimensional tensors, in which case each integer corresponds to one dimension, counting from the leading (leftmost) dimension onward. If the number of indices is smaller than the number of dimensions, a view of the subdimensional tensor is returned.
a = Tensor(torch.tensor([[[[1], [2], [3]], [[4], [5], [6]]],
[[[-1], [-2], [-3]], [[-4], [-5], [-6]]]]), 2)
# Single element indexing along batch dimensions
print("a.batch[1, 0] =")
print(a.batch[1, 0], "\n")
print("a.batch[0] =")
print(a.batch[0], "\n")
# Single element indexing along base dimensions
print("a.base[2, 0] =")
print(a.base[2, 0], "\n")
print("a.base[1] =")
print(a.base[1])
Output:
a.batch[1, 0] =
-1
-2
-3
[ CPULongType{3,1} ]
<Tensor of shape [][3, 1]>
a.batch[0] =
(1,.,.) =
1
2
3
(2,.,.) =
4
5
6
[ CPULongType{2,3,1} ]
<Tensor of shape [2][3, 1]>
a.base[2, 0] =
3 6
-3 -6
[ CPULongType{2,2} ]
<Tensor of shape [2, 2][]>
a.base[1] =
(1,.,.) =
2
5
(2,.,.) =
-2
-5
[ CPULongType{2,2,1} ]
<Tensor of shape [2, 2][1]>
Slicing
NEML2 supports the same slicing rules as in NumPy and PyTorch. A slice takes the form of
start:stop:step
where start, stop, and step are integers representing the starting index, ending index, and the striding of the sliced tensor view. All of these integers are optional when constructing a slice: start default to 0, stop default to (i.e., std::numeric_limits<Size>::max()), and step default to 1. Note that step must be positive.
It is best to learn slicing from examples. Below are equivalent C++ and Python codes applying the same set of slicing operations on the same tensor.
stop default to 'consuming all remaining elements'
a.batch[12::2] =
12
14
16
18
[ CPUDoubleType{4} ]
step default to 1
a.batch[3:6:] =
3
4
5
[ CPUDoubleType{3} ]
Trailing colon(s) can be omitted
a.batch[3:6] =
3
4
5
[ CPUDoubleType{3} ]
a.batch[17:] =
17
18
19
[ CPUDoubleType{3} ]
The default is therefore equivalent to slicing the entire dimension
a.batch[:] =
0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
[ CPUDoubleType{20} ]
from neml2.tensors import Scalar
import torch
torch.set_default_dtype(torch.double)
# Create a tensor with shape (20;)
a0 = Scalar.full(0.0)
a1 = Scalar.full(19.0)
a = Scalar.linspace(a0, a1, 20)
print("Basic syntax: start:stop:step")
print("a.batch[5:17:2] =")
print(a.batch[5:17:2], "\n")
print("Negative start and stop are counted backward")
print("a.batch[-15:-3:2] =")
print(a.batch[-15:-3:2], "\n")
print("start default to 0")
print("a.batch[:17:3] =")
print(a.batch[:17:3], "\n")
print("stop default to 'consuming all remaining elements'")
print("a.batch[12::2] =")
print(a.batch[12::2], "\n")
print("step default to 1")
print("a.batch[3:6:] =")
print(a.batch[3:6:], "\n")
print("Trailing colon(s) can be omitted")
print("a.batch[3:6] =")
print(a.batch[3:6], "\n")
print("a.batch[17:] =")
print(a.batch[17:], "\n")
print("The default is therefore equivalent to slicing the entire dimension")
print("a.batch[:] =")
print(a.batch[:])
Output:
Basic syntax: start:stop:step
a.batch[5:17:2] =
5
7
9
11
13
15
[ CPUDoubleType{6} ]
<Scalar of shape [6][]>
Negative start and stop are counted backward
a.batch[-15:-3:2] =
5
7
9
11
13
15
[ CPUDoubleType{6} ]
<Scalar of shape [6][]>
start default to 0
a.batch[:17:3] =
0
3
6
9
12
15
[ CPUDoubleType{6} ]
<Scalar of shape [6][]>
stop default to 'consuming all remaining elements'
a.batch[12::2] =
12
14
16
18
[ CPUDoubleType{4} ]
<Scalar of shape [4][]>
step default to 1
a.batch[3:6:] =
3
4
5
[ CPUDoubleType{3} ]
<Scalar of shape [3][]>
Trailing colon(s) can be omitted
a.batch[3:6] =
3
4
5
[ CPUDoubleType{3} ]
<Scalar of shape [3][]>
a.batch[17:] =
17
18
19
[ CPUDoubleType{3} ]
<Scalar of shape [3][]>
The default is therefore equivalent to slicing the entire dimension
a.batch[:] =
0
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
[ CPUDoubleType{20} ]
<Scalar of shape [20][]>
Similar to single element indexing, slicing can also be used to index multidimensional tensors. When the number of slices is smaller than the number of dimensions, a view of the subdimensional tensor is returned.
When indexing multidimensional tensors, having to specify the element index or slicing for each dimension would be cumbersome. To simplify multidimensional indexing, some special syntax and notations are reserved for inferring tensor shapes.
NEML2 currently supports two special symbols:
Ellipsis or ... (only available in Python) is equivalent to one or multiple Slice() or : expanding the rest of the dimensions.
None (when used as a index) is equivalent to NumPy's newaxis which unsqueezes a unit-length dimension at the specified place.
Again, the use of these special symbols are best illustrated by examples.
1.13.2