Tensor Indices for AI: Reading Samples, Tokens, and Features

An AI program reads X[1, 2, 0] and returns 0.8. A nearby expression, X[1, 0, 2], returns 0.3. Both addresses are valid. Only one means “the first feature of the last token in the second example.”

The difference is an index: a position used to locate a value inside an ordered structure. Indices are small pieces of notation, but they determine which example, token, feature, pixel, or model output a calculation actually uses. A wrong index can therefore produce a believable number instead of a helpful error.

By the end of this lesson, you will be able to read indexed samples and coordinates, follow a batch-token-feature tensor address, translate a mathematical subscript into NumPy, and detect an address whose numbers are legal but whose axis meanings are wrong. This adds order to the dataset collections from Lesson 4. The official Mathematics for Machine Learning companion places this notation on the route from numerical data to the vectors, matrices, and models used later in the course.

An index answers “which position?”

Suppose an ordered dataset contains three examples:

The subscript in $x_2$ says “select the example at position 2.” It does not say that the example’s value is 2. If the second example is the pair $x_2=(1.7,-0.4)$, then $x_2$ names the whole pair.

Add a second subscript to select a coordinate inside that example:

Read $x_{2,1}$ from left to right: choose example 2, then coordinate 1. The indices form an address; the number stored at that address is the value.

This also separates position from identity. Moving a sample from position 2 to position 7 changes its index even if its stable ID, image pixels, and label do not change. Indices are useful addresses, not permanent names for real-world objects.

Each subscript chooses one axis

An axis is one direction in which an array is organized. A table has a row axis and a column axis. Here, a tensor is a multidimensional array: it can have those two axes or add more axes for concepts such as batches and tokens.

Consider an activation tensor

where:

• $B$ is the number of examples in the batch;
• $T$ is the number of token positions per example; and
• $F$ is the number of features stored for each token.

The notation $\mathbb{R}^{B\times T\times F}$ says that every entry is a real number and that the tensor has shape $B\times T\times F$. One scalar entry is

which means: choose batch position $b$, token position $t$, and feature position $f$, in that order.

For a tensor with shape $2\times3\times4$, a one-based mathematical convention allows

The shape supplies the range of each index. The axis names supply the meaning. Neither fact replaces the other.

One address, worked all the way to a value

Use the following two-batch tensor. Each batch item contains three token rows, and every token has four feature values.

This stepwise reading matters more than memorizing a letter order. Another project might store tokens before batches or use height, width, and channel axes for images. The code and documentation must state the contract.

NumPy starts at zero

Many mathematics texts label positions beginning at 1. Python and NumPy use zero-based indexing, so the first position is 0. NumPy’s official indexing guide also confirms that one integer per dimension selects one element and that the valid nonnegative index for an axis of size $d$ satisfies $0\le i<d$.

Here is the complete tensor and several verified selections:

import numpy as np

X = np.array([
    [[ 0.2,  0.8, -0.1,  0.4],
     [ 1.1, -0.3,  0.5,  0.0],
     [ 0.7,  0.6,  0.2, -0.4]],
    [[-0.2,  0.9,  0.3,  0.1],
     [ 0.4, -0.7,  1.2,  0.5],
     [ 0.8,  0.1, -0.6,  1.5]],
])

print("shape:", X.shape)
print("one activation:", X[1, 2, 0])
print("last token of example 2:", X[1, 2])
print("feature 1 of every last token:", X[:, 2, 0])

shape: (2, 3, 4)
one activation: 0.8
last token of example 2: [ 0.8  0.1 -0.6  1.5]
feature 1 of every last token: [0.7 0.8]

The colon in X[:, 2, 0] means “keep every position on this axis.” The first axis remains, while token and feature receive specific indices, so the result contains one value for each of the two batch items.

The shape attribute reports (2, 3, 4), but the tuple itself does not say “batch, token, feature.” Those names come from the program’s data contract.

A legal address can still express the wrong question

Return to the opening expressions:

X[1, 2, 0]  # 0.8: batch 2, token 3, feature 1
X[1, 0, 2]  # 0.3: batch 2, token 1, feature 3

Both addresses fit the shape. NumPy cannot know that the second line swapped the token and feature meanings. This is a semantic indexing bug: the program retrieves a real entry, but not the entry the calculation intended.

Useful defenses include naming the axis contract beside the tensor, checking the shape at boundaries, and assigning index variables descriptive names:

batch_index = 1
token_index = 2
feature_index = 0

activation = X[batch_index, token_index, feature_index]
assert X.shape == (2, 3, 4)
assert activation == 0.8

The assertion checks this example. The descriptive variables preserve the reasoning that led to the address.

Bounds errors are safer than silent axis errors

For the zero-based shape (2, 3, 4), the last nonnegative address is X[1, 2, 3]. X[2, 0, 0] is invalid because batch index 2 would request a third item from an axis containing only two items. NumPy raises IndexError rather than returning a made-up value.

Negative indices are a Python convention: -1 selects the final position of an axis, so X[-1, -1, -1] is another way to read X[1, 2, 3]. That shorthand is valid code, but it should not be silently translated into a one-based mathematical subscript. State the convention whenever the distinction matters.

The next lesson will use these addresses inside sums. Once $x_i$ clearly means the $i$th example and $\ell_i$ means its loss, a symbol such as $\sum_i\ell_i$ becomes a compact instruction to visit every indexed loss and combine the values.

Check your understanding

activation = X[1, 2, 0]  # 0.8