A tensor is a shape, some strides, and a pointer

Every machine learning library, from NumPy to PyTorch, is built on a single data structure: the tensor. Images, sentences, model weights and gradients are all tensors by the time the maths starts. This post explains what one is, and then shows the surprisingly small piece of code that Ferro uses to represent it.

#What is a tensor?

A tensor is a grid of numbers with any number of dimensions. You already know the first few:

rank 0   a single number            7
rank 1   a list (vector)            [7, 2, 9]
rank 2   a table (matrix)           [[7, 2, 9],
                                     [1, 0, 4]]
rank 3   a stack of tables          two or more matrices, one behind another

The number of dimensions is the tensor's rank, and each dimension is an axis. The shape lists how long the tensor is along each axis: the matrix above has shape [2, 3] — two rows, three columns.

Real data fills these shapes quickly:

• A colour photo is a rank-3 tensor of shape [3, 224, 224]: three colour channels (red, green, blue), each a 224 × 224 grid of pixel brightnesses.
• A batch of 32 such photos, the unit a network actually trains on, is rank 4: [32, 3, 224, 224].
• A batch of sentences becomes [batch, words, features]: each word turned into a list of numbers.

Two rules make a tensor a tensor rather than a nested list. Every element has the same type (all 32-bit floats, say), and the grid is rectangular: every row has the same length. Those rules are what let the numbers sit in one flat block of memory, and that block of memory is where the interesting part begins.

#Ferro's tensor, in full

Here is the entire definition of Ferro's tensor type, minus the documentation:

pub struct Tensor {
    storage: Arc<Storage>,
    shape: Shape,
    strides: Strides,
    offset: usize,
    dtype: DType,
    device: Device,
}

Six fields, and only one of them holds any numbers:

Everything except storage is a description of how to read the numbers. Almost every useful property of an array library falls out of that split, so it is worth going slowly.

#Memory is a line, not a grid

Take a 2 × 3 grid of values:

1  2  3
4  5  6

Computer memory is not a grid. It is one long line of addresses. So the grid has to be flattened, and there are two obvious ways to do it:

row-major:      1  2  3  4  5  6      (walk each row, then move down)
column-major:   1  4  2  5  3  6      (walk each column, then move right)

Neither is more correct. NumPy and C pick row-major; Fortran, MATLAB and most linear algebra libraries pick column-major — which is why you occasionally find yourself transposing a matrix for no reason other than to hand it to somebody else's code.

Once the values are on a line, you need a rule for finding element [i, j]. For the row-major layout it is:

The 3 is there because moving down one row means skipping three elements. Element [1, 2] (second row, third column) sits at position $3 \times 1 + 2 = 5$, and position 5 holds 6. Correct.

That number — how far to step in memory to move one place along an axis — is the stride. A 2 × 3 row-major array has strides (3, 1): step three to change row, one to change column.

For any number of dimensions, the rule becomes:

That formula is the whole of Ferro's Strides::displacement. The formula is not the interesting part. The interesting part is that the strides are data, not a rule baked into the code — which means they can be changed.

#Transposing without moving anything

Transposing swaps rows and columns: the 2 × 3 grid becomes 3 × 2.

original       transpose
1  2  3        1  4
4  5  6        2  5
               3  6

The obvious implementation allocates a new buffer and copies every element to its new home, which takes time proportional to the size of the tensor.

But look at what a transpose actually asks for. Element [i, j] of the transpose is element [j, i] of the original. The original uses strides (3, 1), so reading the same memory with strides (1, 3) gives exactly the transpose. The buffer does not change at all. Only the description does:

let a = Tensor::arange(0..6, DType::F32).reshape([2, 3])?;
let b = a.t()?;

assert_eq!(a.shape(),   &[2, 3]);
assert_eq!(a.strides(), &[3, 1]);
assert_eq!(b.shape(),   &[3, 2]);
assert_eq!(b.strides(), &[1, 3]);

assert!(a.shares_storage_with(&b));    // same buffer, two readings

The transpose takes constant time and allocates nothing, however large the tensor. A tensor that reads another tensor's buffer through a different description is called a view.

The same trick covers a surprising amount of ground:

• Selecting rows 1 and 2 of a matrix: move the offset forward by one row, and shorten the first axis. The strides do not change.
• Dropping a length-1 axis ([1, 3] → [3]): delete one entry from the shape and the matching entry from the strides.
• Reordering axes ([batch, height, width, channels] → [batch, channels, height, width]): reorder both lists the same way.

In Ferro, permute, transpose, t, narrow, slice, squeeze, unsqueeze, expand, broadcast_to and flip are all views. None of them copies a single element.

#What Arc is doing there

If several tensors read the same buffer, who owns the buffer, and when is it freed?

Arc<Storage> is Rust's answer: an atomically reference-counted pointer. The buffer keeps a count of how many Arcs point at it. Cloning an Arc adds one to the count; dropping one subtracts one; when the count reaches zero, the buffer is freed. "Atomically" means the counter uses hardware instructions that are safe when several threads update it at once — which matters, because Ferro's data loader will hand tensors between threads.

The consequence worth internalising:

let a = Tensor::zeros([1000, 1000], DType::F32);
let b = a.clone();

b is not a copy of a million floats. It is a second pointer to the same million floats, and creating it took about as long as incrementing an integer — because that is what it did.

#Two decisions worth explaining

Strides are signed. They are isize, not usize, so they can be negative. A negative stride means "walk backwards through memory", and that is exactly what reversing an axis is:

buffer:            10  20  30  40
original:          offset 0, stride +1   →  10  20  30  40
flipped (a view):  offset 3, stride −1   →  40  30  20  10

Point the offset at the last element, step backwards, and flip joins the list of operations that cost nothing. Supporting negative strides later would have meant revisiting every stride calculation in the codebase once there were dozens of them. Supporting them from the start cost one type choice.

Rank is decided at run time. A Ferro tensor does not know its number of dimensions at compile time. Rust could express that — Tensor<f32, 4> for a rank-4 tensor of floats — and the compiler would then catch shape mistakes before the program ran, which is genuinely appealing.

It was rejected for two practical reasons. A library that loads arbitrary model weights cannot know their ranks in advance. And a list of a network's parameters holds tensors of several different ranks at once. Both are unavoidable, and both are painful when the rank is part of the type. A typed wrapper can be built on top of an untyped tensor later; going the other way is much harder.

#Where the numbers live

Storage is an enum with one variant per element type:

pub enum Storage {
    Bool(Vec<bool>),
    U8(Vec<u8>),
    I32(Vec<i32>),
    I64(Vec<i64>),
    F32(Vec<f32>),
    F64(Vec<f64>),
}

The fastest design is something more sophisticated: a single byte buffer, aligned to 64 bytes so that SIMD instructions — which process several numbers in one instruction — can load from it efficiently. Ferro will get there, but not yet, and the reasoning is worth stating plainly because it is a deliberate choice.

Reinterpreting a buffer of bytes as a buffer of floats requires unsafe in Rust: you are asserting something the compiler cannot check, namely that the bytes are correctly aligned and sized for the type you claim. That is a perfectly reasonable thing to do, in the right place.

The right place is not here. Right now the job is getting index arithmetic exactly right, and index bugs and memory-safety bugs look identical from the outside: both produce wrong numbers. Doing both at once would give every wrong answer two possible explanations. With typed vectors, an out-of-range index is a clean panic with a line number, and the tensor code contains no unsafe at all — so Miri, the tool that checks unsafe code for undefined behaviour, has a clean baseline to compare against later.

Storage sits behind an Arc, and nothing outside its own file looks inside it. Swapping in an aligned allocator alongside the SIMD work that needs it will be a change to one file.

#The check that says it works

One expression serves as the acceptance test for all of this:

let a = Tensor::arange(0..24, DType::F32).reshape([2, 3, 4])?;
let b = a.permute(&[2, 0, 1])?.narrow(0, 1, 2)?.contiguous();
assert_eq!(b.shape(), &[2, 2, 3]);

It is a good test because it exercises the whole chain. Here is what each step does to the description:

The permute produces strides that are no longer in decreasing order. The narrow moves the offset without touching the strides, so the tensor now starts partway into its buffer. Only contiguous touches the data: it walks that unusual layout in logical order and writes the elements out packed, into a new buffer:

[[[ 1,  5,  9],
  [13, 17, 21]],

 [[ 2,  6, 10],
  [14, 18, 22]]]

You can check any entry by hand. b[1, 0, 2] sits at offset $1 + 1 \times 1 + 0 \times 12 + 2 \times 4 = 10$ in the original buffer, and arange(0..24) stores the value 10 there.

The expected output was generated by NumPy and committed before any of this code existed. When the implementation landed, the assertion did not change. That is the payoff of writing the test first: the target was fixed before the implementation had a chance to define its own idea of correct.

#The idea to keep

A tensor is a flat buffer of numbers plus a small description — shape, strides and offset — of how to read it. Change the description and you get a new tensor for free. Copy the numbers only when you have to.

Next in the series: stride tricks — broadcasting, expanding, and what else becomes possible by bending the strides.