PyTorch tensors

The same operations, with dim instead of axis

Most of what you did with NumPy arrays works the same way on tensors. Run this and compare each line with what NumPy would do:

import numpy as np
import torch

x = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
print(x * 2 + 1)                        # element-wise, like NumPy
print(x @ x.T)                          # matrix multiply: (2, 3) @ (3, 2)
print(x.sum(dim=0), x.mean(dim=1))      # reductions take dim= (NumPy: axis=)
print(x.max(dim=1))                     # max over a dim returns values AND indices
print(x + torch.tensor([10.0, 20.0, 30.0]))   # broadcasting works the same way
print(x[:, 1], x[x > 3])                # slicing and boolean masks too

n = x.numpy()                           # to NumPy (shares memory)...
t = torch.from_numpy(np.ones(3))        # ...and back
print(type(n).__name__, t.dtype)

Line by line:

• x * 2 + 1 works on every number on its own. 1 becomes 3, 2 becomes 5, and so on.
• x @ x.T is matrix multiplication. x.T is x flipped, so its shape is (3, 2). A (2, 3) times a (3, 2) gives a (2, 2) result. The top-left number is 1×1 + 2×2 + 3×3 = 14.
• x.sum(dim=0) adds down each column: 1 + 4 = 5, 2 + 5 = 7, 3 + 6 = 9.
• x.mean(dim=1) averages along each row: the mean of 1, 2, 3 is 2, and of 4, 5, 6 is 5.
• x.max(dim=1) finds the biggest number in each row. It returns two things: the values (3 and 6), and the positions where they sit (both at index 2).
• Adding a 3-number tensor adds it to every row. That's broadcasting, as in NumPy.
• x[:, 1] takes column 1 from every row. x[x > 3] keeps only the numbers above 3.

The word reduction means any operation that shrinks a dimension away, such as sum, mean or max.

In plain words: dim names the dimension you squash. dim=0 squashes the rows, so you get one answer per column. dim=1 squashes the columns, so you get one answer per row.

Two differences catch people. First, reductions use dim where NumPy uses axis. Second, max(dim=...) returns a pair (values, indices), not only the values. The indices are useful: they tell you which item scored highest, such as a model's top-scoring answer.

The last three lines move data between the two libraries. x.numpy() turns a tensor into a NumPy array. torch.from_numpy goes the other way. On the CPU, both share memory: they're two views of the same numbers, so changing one changes the other. Notice that the new tensor kept NumPy's dtype, float64. PyTorch doesn't quietly convert it.

Reshaping

Neural-network code spends a lot of time moving dimensions around. Picture an egg carton of 24 eggs. You can see it as 2 rows of 12, or 4 rows of 6. The eggs don't move; only the way you group them changes. Reshaping a tensor works the same way. Here are the common moves:

import torch

x = torch.arange(24)
print(x.view(2, 3, 4).shape)            # same 24 numbers seen as 2 × 3 × 4
print(x.view(4, -1).shape)              # -1: "work this one out"
b = x.view(2, 3, 4)
print(b.transpose(1, 2).shape)          # swap two dims: (2, 4, 3)
print(b.permute(2, 0, 1).shape)         # reorder all dims: (4, 2, 3)
print(b[0].unsqueeze(0).shape)          # add a dim of size 1: (1, 3, 4)
print(b.flatten(1).shape)               # merge everything after dim 0: (2, 12)

batch = torch.randn(32, 10, 64)         # (batch, sequence, features): the shape LLMs pass around
print((batch @ torch.randn(64, 16)).shape)   # matmul applies to the last dim

What each move does:

• view(2, 3, 4) regroups the same 24 numbers into 2 blocks of 3 rows of 4. The sizes must multiply to the total: 2 × 3 × 4 = 24.
• view(4, -1) asks for 4 rows. The -1 means "work this one out". 24 ÷ 4 = 6, so the shape is (4, 6).
• transpose(1, 2) swaps two dimensions. (2, 3, 4) becomes (2, 4, 3).
• permute(2, 0, 1) lists a new order for all the dimensions. Old dim 2 (size 4) goes first, then old dim 0 (size 2), then old dim 1 (size 3). That gives (4, 2, 3).
• unsqueeze(0) adds a new dimension of size 1 at the front. It turns one example into a "batch of one", which models expect.
• flatten(1) keeps dim 0 and merges everything after it. 3 × 4 = 12, so (2, 3, 4) becomes (2, 12).

In plain words: view and flatten regroup the numbers. transpose and permute change the order of the dimensions. unsqueeze adds a size-1 dimension. None of them change how many numbers there are.

The shape inside a language model

The last two lines show the layout you'll see most: (batch, sequence length, features).

• batch: how many examples go through together. Here, 32 sequences of text.
• sequence length: how many tokens each sequence has. Here, 10.
• features: how many numbers describe each token. Here, 64.

So the model keeps one row of 64 features per token, and one sequence per batch item.

Multiplying by a 2-D weight with shape (64, 16) acts on the last dimension only. Every token's 64 features become 16 new features. PyTorch does this for all 32 × 10 = 320 tokens at once. The result has shape (32, 10, 16).

import torch
x = torch.arange(12)
print(x.view(3, -1).shape)

import torch

h = torch.arange(24.0).view(2, 3, 4)   # 2 messages, 3 tokens each, 4 features per token
print(h.shape)
print(h[0])                            # message 0: one row per token
print(h.mean(dim=1))                   # average the tokens: one row per message
w = torch.ones(4, 2)                   # a tiny weight: 4 features in, 2 out
out = h @ w
print(out.shape)
print(out[0])

Devices and dtypes

A GPU is like a kitchen with a thousand cooks. Each cook is slow, but they all chop at once. A neural network is mostly the same simple sums repeated millions of times, so a GPU suits it well.

On a machine with an NVIDIA GPU, moving work there is one call. (Those GPUs go by the name "cuda", after NVIDIA's software.)

device = "cuda" if torch.cuda.is_available() else "cpu"
x = torch.randn(1000, 1000, device=device)     # created on the GPU
model = model.to(device)                       # move a model's weights
y = model(x)                                   # runs on the GPU; inputs must be on the same device

The first line picks the GPU if there is one, and the CPU otherwise. The second creates a tensor directly on that device. The third moves a model's weights there with .to(device). The fourth runs the model, and it works because x and the weights now sit on the same device.

Here's the catch. Tensors on different devices can't be combined. It's like asking one cook to mix two bowls that sit in different kitchens. Try it and you get the classic error "Expected all tensors to be on the same device". The fix is to move the inputs (or the model) with .to(device) so they match.

On a Mac with Apple silicon, the device is "mps" instead of "cuda".

Dtypes matter most on a GPU, because GPU memory is limited and costly. Lower-precision dtypes, float16 and bfloat16, use 2 bytes per number. That halves memory again compared with float32. Large models are trained and served this way.

Some quick sums show why. A model with 7 billion weights needs 7 billion × 4 bytes = 28 GB in float32. In bfloat16 it needs 7 billion × 2 bytes = 14 GB. The GPUs lesson covers the rest.

In this browser subset, everything stays on the CPU. torch.cuda.is_available() returns False.