15 min

The tensor

Everything in ML is arithmetic on arrays of numbers.

Milestone a+b, a*b and a@b compute correctly.

PyTorch is called PyTorch because of the tensor. Everything that flows through a neural network is one: the input data, the weights, every intermediate result and, later on, every gradient. And a tensor is nothing mysterious - it is a field of numbers with a shape. A greyscale image is a tensor of shape (28, 28), a batch of 32 such images one of shape (32, 28, 28).

This first step builds exactly that shell. It cannot know anything about gradients yet, that comes in the next chapter. It can do three things: add, multiply elementwise and take the matrix product. A forward pass needs nothing more.

The goal

Everything in ML is arithmetic on arrays of numbers. Saying that out loud once removes a good deal of the reverence around the subject. A network does not "understand" anything. It multiplies one matrix by another, adds a vector and cuts off negative values. The three operations in this chapter are already almost all of what happens in a forward pass.

What matters is the difference between the two multiplications, because it is the most common source of errors later on:

  • a * b is elementwise: same shape in, same shape out. Every number meets its partner at the same position.
  • a @ b is the matrix product: (n, k) times (k, m) gives (n, m). Here dimensions collapse into each other, and that is precisely what a layer does.

Zero magic in this step

The rule of the course applies from the very first line: NumPy is allowed, torch, jax, tinygrad and autograd are forbidden. NumPy takes the arithmetic on fields of numbers off your hands - np.asarray, +, *, @ - and that is fine, because fast array arithmetic is not the point of this course. The point is the shell around it, and you write that yourself.

The scaffold

Set up your repo. Two levels, because meintorch is both the project and the package inside it - the usual layout for Python packages:

meintorch/                  ← your repo
├── meintorch/              ← the package, your framework
│   ├── __init__.py         (empty)
│   └── tensor.py
└── tests/
    └── test_01_tensor.py

Always run the tests from the root of the repo with python -m pytest. The -m is not a matter of taste: it puts the current directory on the import path, and that is the only reason from meintorch.tensor import Tensor finds anything at all.

The exercise

This is coach mode: you get the scaffold and the gaps, you write the core. meintorch/tensor.py:

import numpy as np


def as_tensor(x):
    """Numbers and lists become tensors, tensors stay what they are."""
    return x if isinstance(x, Tensor) else Tensor(x)


class Tensor:
    """A field of numbers with a shape. At this point it is nothing more."""

    def __init__(self, data):
        self.data = np.asarray(data, dtype=np.float64)

    @property
    def shape(self):
        return self.data.shape

    def __repr__(self):
        return f"Tensor({self.data})"

    def __add__(self, other):
        # TODO(du): elementwise addition, result wrapped in a Tensor again
        raise NotImplementedError

    def __mul__(self, other):
        # TODO(du): elementwise multiplication - not the matrix product!
        raise NotImplementedError

    def __matmul__(self, other):
        # TODO(du): the matrix product
        raise NotImplementedError

Hint 1 - the direction

Each of the three methods is a three-liner following the same pattern: turn the other side into a tensor, let NumPy do the arithmetic, wrap the result in a tensor again. None of it needs a loop.

Hint 2 - more concrete

self.data and other.data are NumPy arrays. NumPy already knows the right operator for all three cases; you are only passing it on. And as_tensor is not sitting at the top by accident - without that call a * 2.0 fails, because a number has no .data.

Show the solution
    def __add__(self, other):
        other = as_tensor(other)
        return Tensor(self.data + other.data)

    def __mul__(self, other):
        other = as_tensor(other)
        return Tensor(self.data * other.data)

    def __matmul__(self, other):
        other = as_tensor(other)
        return Tensor(self.data @ other.data)

In goal mode

The same step, autonomously: the agent sets up repo and package, implements the three operators, writes the milestone test along with them and runs it until it is green. It does not ask, it hands over a finished chapter. The instruction for it:

Build prompt
Set up a Python repo `meintorch/` with the package `meintorch/meintorch/`.
In `meintorch/tensor.py`, implement a class `Tensor` that wraps a NumPy array
(dtype float64) and supports `__add__`, `__mul__` and `__matmul__` - elementwise
addition, elementwise multiplication, matrix product. Scalars and lists on the
right-hand side have to work as well.
Mandatory: NumPy is allowed, torch/jax/tinygrad/autograd are forbidden.
Write `tests/test_01_tensor.py` alongside it and run `python -m pytest` until
every test is green. Show me the test output at the end.

The milestone test

tests/test_01_tensor.py. Green means the chapter is done.

import numpy as np

from meintorch.tensor import Tensor


def test_addition_is_elementwise():
    a = Tensor([1.0, 2.0, 3.0])
    b = Tensor([10.0, 20.0, 30.0])
    assert np.allclose((a + b).data, [11.0, 22.0, 33.0])


def test_multiplication_is_elementwise():
    a = Tensor([1.0, 2.0, 3.0])
    b = Tensor([10.0, 20.0, 30.0])
    assert np.allclose((a * b).data, [10.0, 40.0, 90.0])


def test_matmul_is_the_matrix_product():
    a = Tensor([[1.0, 2.0], [3.0, 4.0]])
    b = Tensor([[5.0, 6.0], [7.0, 8.0]])
    assert np.allclose((a @ b).data, [[19.0, 22.0], [43.0, 50.0]])


def test_shape_and_scalars():
    a = Tensor([[1.0, 2.0, 3.0]])
    assert a.shape == (1, 3)
    assert np.allclose((a * 2.0).data, [[2.0, 4.0, 6.0]])