Math & Python Prereqs

The math you need for the first 14 chapters of obvix-learn is the math you need for one specific operation. That operation is gradient descent on a parameterized function evaluated against data. Every classical ML algorithm, every neural network, every transformer, every diffusion model, every RLHF policy, every mech-interp probe, is some variant of:

1. Define a function $f_\theta: \mathcal{X} \to \mathcal{Y}$ parameterized by $\theta$.
2. Define a loss $L(\theta; \mathcal{D})$ that scores $f_\theta$ on data $\mathcal{D}$.
3. Compute $\nabla_\theta L$.
4. Update $\theta \leftarrow \theta - \eta \nabla_\theta L$.

That's it. The reason you need linear algebra is that $\theta$ is almost always a collection of matrices, and $\nabla_\theta L$ is a matrix-shaped object you compute by chain rule. The reason you need probability is that the loss $L$ usually has the form $-\log p_\theta(y \mid x)$, a negative log-likelihood. The reason you need information theory is to know why cross-entropy is the right loss to minimize when you don't know the true data distribution.

Everything in this chapter exists to make those four bullet points readable.

You have two reasonable options for running every code block in this curriculum: Google Colab (zero install, free GPU access, browser-based) or local Python (more control, faster iteration, no quota limits). Pick Colab for the first ten chapters. Pick local for the rest.

Colab path. Go to colab.research.google.com. Create a new notebook. Type the following in the first cell and run it:

import sys, torch, numpy as np, pandas as pd, sklearn
print(f"Python {sys.version.split()[0]}")
print(f"PyTorch {torch.__version__}, CUDA available: {torch.cuda.is_available()}")
print(f"NumPy {np.__version__}, pandas {pd.__version__}, sklearn {sklearn.__version__}")

You should see Python 3.10+, PyToranother chapter.x, CUDA True if you selected the T4 runtime. Done.

Local path (Linux/macOS; Windows users add wsl or use Anaconda):

# 1. Create a project venv (use Python 3.10 or newer)
python3 -m venv ~/.venvs/obvix
source ~/.venvs/obvix/bin/activate

# 2. Upgrade pip
pip install --upgrade pip

# 3. Install the core stack
pip install numpy pandas matplotlib scikit-learn jupyter

# 4. PyTorch — pick the right index URL for your CUDA version
#    For CPU only:
pip install torch --index-url https://download.pytorch.org/whl/cpu
#    For CUDA 12.1:
# pip install torch --index-url https://download.pytorch.org/whl/cu121

# 5. Test
python -c "import torch; print(torch.__version__, torch.cuda.is_available())"

If you want notebooks locally: jupyter lab. If you want a tighter editor: VS Code with the Python extension opens .ipynb files natively. If you want package isolation per project: uv from Astral is the fastest tool and the one I use now.

NumPy is the array library every other Python ML library is built on. The only object you need to know is np.ndarray. It is a multi-dimensional, homogeneous, typed buffer with a shape and a stride.

import numpy as np

# Construction
a = np.array([1, 2, 3, 4])              # shape (4,), dtype int64
b = np.zeros((3, 4))                    # shape (3, 4), dtype float64
c = np.ones((2, 3), dtype=np.float32)   # shape (2, 3), dtype float32
d = np.arange(0, 10, 2)                 # [0, 2, 4, 6, 8]
e = np.linspace(0, 1, 5)                # [0., 0.25, 0.5, 0.75, 1.]
r = np.random.randn(3, 4)               # standard normal, shape (3, 4)

# Shape and dtype
r.shape       # (3, 4)
r.dtype       # float64
r.ndim        # 2
r.size        # 12

# Indexing and slicing
r[0]          # first row, shape (4,)
r[:, 0]       # first column, shape (3,)
r[1:3, :2]    # rows 1-2, cols 0-1, shape (2, 2)
r[r > 0]      # boolean mask, returns 1D array of all positives

# Reshape (no copy, just stride change)
r.reshape(4, 3)
r.reshape(-1)         # flatten to 1D
r.T                   # transpose, shape (4, 3)

# Arithmetic: element-wise by default
a + 1                 # broadcast scalar
r + r                 # element-wise
r * 2                 # element-wise
r ** 2                # element-wise

# Matrix multiplication: explicit
A = np.random.randn(3, 4)
B = np.random.randn(4, 5)
C = A @ B             # shape (3, 5)
# Equivalent: np.matmul(A, B), np.dot(A, B) for 2D

# Reductions
r.sum()               # scalar
r.sum(axis=0)         # shape (4,) — sum each column
r.sum(axis=1)         # shape (3,) — sum each row
r.mean(), r.std(), r.max(), r.argmax()

The two ideas that trip up beginners: broadcasting and the axis parameter.

Broadcasting. When you add a (3, 4) array to a (4,) array, NumPy stretches the second to (3, 4) by repeating along axis 0. The full rule: dimensions of size 1 (or absent) are stretched to match. (3, 1) + (1, 4) = (3, 4). This is the operation that makes vectorized code short. It is also the operation that produces silent shape bugs when you forget that (3,) and (3, 1) are different.

Axis. r.sum(axis=0) means "collapse axis 0", returning an array with one fewer dim along that axis. Get this wrong and your loss curve looks fine but the gradient is for the wrong quantity. The fastest debugging trick in machine learning: print(x.shape) everywhere.

The 800-page linear algebra textbook covers eight semesters of material. Machine learning uses about four operations from it. Here they are.

One piece of notation first, since it runs through everything below: $\mathbb{R}^n$ means a vector of $n$ real numbers (a list of length $n$), and $\mathbb{R}^{m \times n}$ means a matrix with $m$ rows and $n$ columns. The shape is the whole point — reading "$A \in \mathbb{R}^{m \times n}$" as "$A$ is $m$-by-$n$" is most of what you need.

Vector addition and scalar multiplication

For $\mathbf{x}, \mathbf{y} \in \mathbb{R}^n$ and $\alpha \in \mathbb{R}$:

Element-wise. This is gradient descent's update step: $\mathbf{w}_{t+1} = \mathbf{w}_t - \eta \mathbf{g}_t$.

Inner (dot) product

For $\mathbf{x}, \mathbf{y} \in \mathbb{R}^n$:

The most-used operation in ML. A neuron's pre-activation is $\mathbf{w}^\top \mathbf{x} + b$. An attention score is $\mathbf{q}^\top \mathbf{k}$. Cosine similarity is $\frac{\mathbf{x}^\top \mathbf{y}}{\|\mathbf{x}\|\|\mathbf{y}\|}$.

x = np.array([1., 2., 3.])
y = np.array([4., 5., 6.])
x @ y          # 32.0
np.dot(x, y)   # same
(x * y).sum()  # same, slower path

Matrix-vector and matrix-matrix product

For $A \in \mathbb{R}^{m \times n}$, $\mathbf{x} \in \mathbb{R}^n$:

For $A \in \mathbb{R}^{m \times k}$, $B \in \mathbb{R}^{k \times n}$:

The thing to remember: inner dimensions must match, outer dimensions become the output shape. (3, 4) @ (4, 5) = (3, 5). (3, 4) @ (5, 4) is an error.

When two shapes almost line up but the inner dimensions are swapped, the fix is the transpose: $A^\top$ (written A.T in NumPy) flips a matrix's rows and columns, so an $m \times n$ matrix becomes $n \times m$ with $A^\top_{ij} = A_{ji}$. It is the standard way to rotate one operand so the inner dimensions meet.

A fully-connected layer is one matrix-matrix product plus a broadcasted vector add:

where $X \in \mathbb{R}^{B \times d_{\text{in}}}$, $W \in \mathbb{R}^{d_{\text{out}} \times d_{\text{in}}}$, $\mathbf{b} \in \mathbb{R}^{d_{\text{out}}}$, $\mathbf{Y} \in \mathbb{R}^{B \times d_{\text{out}}}$.

B, d_in, d_out = 32, 100, 50
X = np.random.randn(B, d_in)
W = np.random.randn(d_out, d_in)
b = np.random.randn(d_out)

Y = X @ W.T + b      # (32, 100) @ (100, 50) + (50,) -> (32, 50)
assert Y.shape == (B, d_out)

Norms

The $L^2$ (Euclidean) norm of a vector:

The $L^1$ norm (Manhattan):

The $L^p$ norm in general:

You use $L^2$ for distances and energies, $L^1$ for sparsity penalties, $L^\infty = \max_i |x_i|$ for worst-case bounds. The Frobenius norm of a matrix is $\|A\|_F = \sqrt{\sum_{ij} A_{ij}^2}$. It is the $L^2$ norm of the flattened matrix.

np.linalg.norm(x)           # L2
np.linalg.norm(x, ord=1)    # L1
np.linalg.norm(x, ord=np.inf)  # Linf
np.linalg.norm(A)           # Frobenius for 2D

Two operations to know about but not memorize.

Inverse and pseudo-inverse. For square non-singular $A$, the inverse $A^{-1}$ satisfies $AA^{-1} = I$. You almost never compute it explicitly in ML; you solve $A\mathbf{x} = \mathbf{b}$ via np.linalg.solve(A, b) instead. For non-square $A$, the Moore-Penrose pseudo-inverse $A^+$ generalizes this. It shows up in the closed-form solution to least squares: $\hat{\mathbf{w}} = (X^\top X)^{-1} X^\top \mathbf{y}$, which you basically never compute that way in practice (you use SGD or sklearn.linear_model.LinearRegression).

Eigendecomposition and SVD. $A = Q \Lambda Q^{-1}$ when $A$ is symmetric. $A = U \Sigma V^\top$ for any matrix. These power dimensionality reduction (PCA), and they show up in transformer training as the singular value spectrum of weight matrices. another chapter is the chapter where they earn their keep.

ML uses three calculus tools: scalar derivatives, the chain rule, and vector gradients. Everything else (Taylor series, integration by parts, the implicit function theorem) lives outside the curriculum.

Scalar derivatives

The rules you need:

The sigmoid $\sigma(x) = \frac{1}{1 + e^{-x}}$ is the squashing function that maps any real number into $(0, 1)$, which is why it turns a score into a probability for binary classification. Its derivative shows up so often it's worth memorizing in that factored form — and note that the factor $\sigma(x)(1 - \sigma(x))$ goes to zero as the output saturates toward 0 or 1, the fact that section 7 leans on to explain why squared error stalls there but cross-entropy does not.

Chain rule

If $y = f(u)$ and $u = g(x)$, then

This is the rule. Backpropagation is the chain rule applied carefully. The careful part is that you have to remember what activations you saw on the way in.

Example. $y = (3x + 2)^4$. Let $u = 3x + 2$. Then $\frac{dy}{du} = 4u^3 = 4(3x+2)^3$ and $\frac{du}{dx} = 3$, so $\frac{dy}{dx} = 12(3x+2)^3$.

Partial derivatives and gradients

If $f: \mathbb{R}^n \to \mathbb{R}$, the partial derivative with respect to $x_i$ holds all other coordinates fixed:

The gradient is the vector of all partials:

Three facts about the gradient that you need:

1. Direction of steepest ascent. $\nabla f$ points in the direction $f$ increases fastest, locally.
2. Zero at minima. If $\mathbf{x}^*$ is a local minimum and $f$ is smooth, $\nabla f(\mathbf{x}^*) = \mathbf{0}$.
3. Linear in $f$. $\nabla (f + g) = \nabla f + \nabla g$. $\nabla (\alpha f) = \alpha \nabla f$.

Common gradients in ML

Each identity below is derived the same mechanical way: write the scalar function out as a sum over coordinates, take the ordinary partial $\partial / \partial w_i$ holding the rest fixed, then stack those $n$ partials back into a vector — that stacked vector is the gradient. You never need a new rule; it is the scalar chain rule applied one coordinate at a time. The four results worth memorizing, so you don't redo the bookkeeping each time:

For $\mathbf{w} \in \mathbb{R}^n$, $\mathbf{a} \in \mathbb{R}^n$, $A \in \mathbb{R}^{n \times n}$:

The last one is the gradient of the least-squares loss. Setting it to zero gives the normal equations $X^\top X \mathbf{w} = X^\top \mathbf{y}$.

Verifying gradients with autograd

When in doubt, check your hand-derived gradient against PyTorch.

import torch

# Define a simple quadratic
def f(w, X, y):
    return ((X @ w - y) ** 2).sum() / 2

# Random data
torch.manual_seed(0)
X = torch.randn(10, 3)
y = torch.randn(10)
w = torch.randn(3, requires_grad=True)

# Forward + backward
loss = f(w, X, y)
loss.backward()

# Hand-derived gradient: X^T (Xw - y)
manual = X.T @ (X @ w - y)
print("autograd:", w.grad)
print("manual:  ", manual)
assert torch.allclose(w.grad, manual, atol=1e-5)

If assert fails, your math is wrong. If it passes, you can trust the autograd.

Probability shows up in three places: defining loss functions, modeling stochastic optimizers, and reasoning about sampling.

Random variables and distributions

A random variable $X$ is a function from a sample space $\Omega$ to a value space. We almost never write the sample space. What we work with is the distribution $p(X)$.

Discrete: $p(X = x) \geq 0$ and $\sum_x p(X = x) = 1$.

Continuous: $p(x) \geq 0$ (a density, can be $>1$) and $\int p(x)\, dx = 1$. Probabilities are integrals over intervals.

The distributions you'll meet in the first ten chapters:

Expectation

The expected value of $X$ under $p$:

• Discrete: $E[X] = \sum_x x \cdot p(x)$
• Continuous: $E[X] = \int x \cdot p(x)\, dx$

Linearity. $E[X + Y] = E[X] + E[Y]$ always. $E[\alpha X] = \alpha E[X]$. This is the one rule you'll use ten times per derivation.

Variance. $\text{Var}(X) = E[(X - E[X])^2] = E[X^2] - E[X]^2$.

Joint, conditional, marginal

Given a joint $p(X, Y)$:

• Marginal: $p(X) = \sum_y p(X, Y = y)$ (or integrate for continuous)
• Conditional: $p(Y \mid X) = \frac{p(X, Y)}{p(X)}$
• Independence: $X \perp Y$ iff $p(X, Y) = p(X) p(Y)$

Bayes' rule

Reading: posterior = likelihood × prior / evidence. The evidence $p(X) = \sum_y p(X \mid Y = y) p(Y = y)$ is just the normalizer.

Bayes' rule is the basis of every probabilistic classifier you'll meet: naive Bayes, Gaussian discriminant analysis, logistic regression's probabilistic interpretation, Bayesian neural networks. It also justifies why we train classifiers to output $p(Y \mid X)$ and pick $\arg\max_y p(Y = y \mid X)$ as the prediction.

Maximum likelihood: the principle behind every loss

Given data $\mathcal{D} = \{(x_i, y_i)\}_{i=1}^N$ and a parameterized model $p_\theta(y \mid x)$, the maximum likelihood estimate is:

Products of probabilities underflow. We take logs:

That last expression is negative log-likelihood, the most common loss function in deep learning. For Gaussian $p_\theta$ it becomes mean squared error. For categorical $p_\theta$ it becomes cross-entropy. Same idea, different distributional assumption.

You need three quantities. All are measured in nats (natural log) or bits (log base 2). Use nats; PyTorch does.

Entropy

The entropy of a distribution $p$:

Entropy measures uncertainty. A deterministic distribution ($p$ is a one-hot) has $H = 0$. A uniform over $K$ outcomes has $H = \log K$. Intuitively: how surprised you'd be, on average, by a sample.

Cross-entropy

The cross-entropy between $p$ (truth) and $q$ (your model):

This is what you minimize when you train a classifier. If $p$ is the empirical distribution (a one-hot for each training example) and $q = q_\theta(y \mid x)$ is your model's predictive distribution, then cross-entropy is the same loss as negative log-likelihood:

Read that again. Cross-entropy and NLL are two names for the same thing when the target distribution is a one-hot.

KL divergence

The Kullback-Leibler divergence from $q$ to $p$:

KL is non-negative, zero iff $p = q$, not symmetric ($D_{KL}(p \| q) \neq D_{KL}(q \| p)$). It appears as the "distance" between distributions in VAEs, in RLHF, in mech-interp logit-difference analyses.

Why cross-entropy works for classification

Two reasons.

First, it has the right gradient. For softmax output $q_\theta(y = k \mid x) = \frac{e^{z_k}}{\sum_j e^{z_j}}$ and one-hot target $y = k^*$, the gradient of $-\log q_\theta(k^* \mid x)$ with respect to the logits $z$ — the raw, un-normalized scores the model emits just before the softmax turns them into probabilities — is exactly $q - y$. That is the simplest possible gradient: model minus target. It makes training stable.

Second, it correctly penalizes overconfident wrong predictions. Mean squared error on probabilities saturates the gradient when the prediction is near 0 or 1. Cross-entropy does not. A confident wrong prediction (say $q = 0.001$ when the truth is $y = 1$) produces a large loss and a large gradient.

import torch
import torch.nn.functional as F

# logits: (batch, n_classes), targets: (batch,) integer labels
logits = torch.randn(4, 3)
targets = torch.tensor([0, 2, 1, 1])
loss = F.cross_entropy(logits, targets)
print(loss.item())   # one scalar

# What it does under the hood: log_softmax + nll_loss
log_probs = F.log_softmax(logits, dim=-1)
manual = -log_probs[torch.arange(4), targets].mean()
assert torch.allclose(loss, manual)

Every concept from sections 4–7 shows up in one place: deriving and implementing linear regression. This is the smallest end-to-end ML algorithm and it earns its place as the first thing you implement.

Problem. Given $\{(x_i, y_i)\}_{i=1}^N$ with $x_i \in \mathbb{R}^d$ and $y_i \in \mathbb{R}$, find $\mathbf{w} \in \mathbb{R}^d$, $b \in \mathbb{R}$ that minimize:

Probabilistic interpretation. Assume $y_i = \mathbf{w}^\top x_i + b + \epsilon_i$ with $\epsilon_i \sim \mathcal{N}(0, \sigma^2)$. The MLE of $\mathbf{w}, b$ under this model is exactly the minimizer of the squared loss. That is why squared loss is the canonical regression loss: it is NLL under Gaussian noise.

Gradient. Stack the data as $X \in \mathbb{R}^{N \times d}$, $\mathbf{y} \in \mathbb{R}^N$. Absorb $b$ into $\mathbf{w}$ by adding a column of ones to $X$. Then $L = \frac{1}{2N} \|X\mathbf{w} - \mathbf{y}\|^2$ and from section 4:

The same fit done two ways against the library. First, the closed-form solver: sklearn's LinearRegression solves the normal equations directly, where the from-scratch path descends to the same minimum by SGD.

model = LinearRegression()
model.fit(X, y)
print(model.coef_, model.intercept_)

def gradient_descent(X, y, lr=0.1, n_steps=500):
    N, d = X.shape
    w = np.zeros(d)
    b = 0.0
    losses = []
    for step in range(n_steps):
        losses.append(loss(w, b, X, y))
        residual = X @ w + b - y
        dw = X.T @ residual / N
        db = float(residual.mean())
        w = w - lr * dw
        b = b - lr * db
    return w, b, losses

And the PyTorch version of the same descent — nn.Linear trained with torch.optim.SGD — which mirrors the from-scratch loop step for step, just with autograd doing the gradient bookkeeping.

model = nn.Linear(in_features=1, out_features=1)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
loss_fn = nn.MSELoss()
for step in range(500):
    optimizer.zero_grad()
    loss = loss_fn(model(X_t).squeeze(), y_t)
    loss.backward()
    optimizer.step()

def gradient(w, b, X, y):
    N = X.shape[0]
    residual = X @ w + b - y
    dw = X.T @ residual / N
    db = float(residual.mean())
    return dw, db

def gradient_descent(X, y, lr=0.1, n_steps=500):
    N, d = X.shape
    w = np.zeros(d)
    b = 0.0
    losses = []
    for step in range(n_steps):
        losses.append(loss(w, b, X, y))
        dw, db = gradient(w, b, X, y)
        w = w - lr * dw
        b = b - lr * db
    return w, b, losses

Three implementations, three different APIs, same math, same answer. You now have all the pieces of every classical ML algorithm in this curriculum.

NumPy and PyTorch operations dispatch to vectorized C and CUDA kernels. Python for loops do not. The performance gap is two to three orders of magnitude.

import numpy as np
import time

N = 1_000_000
x = np.random.randn(N)

# Slow: Python loop
t0 = time.time()
total_loop = 0.0
for xi in x:
    total_loop += xi ** 2
t_loop = time.time() - t0

# Fast: NumPy vectorized
t0 = time.time()
total_np = (x ** 2).sum()
t_np = time.time() - t0

print(f"Python loop: {t_loop:.3f}s")
print(f"NumPy:       {t_np:.6f}s")
print(f"Speedup:     {t_loop / t_np:.0f}x")

On a typical laptop, the speedup ranges from tens to a few hundred times, depending on machine and array size — it is not a fixed number. On GPU with PyTorch, the speedup over CPU NumPy is another 10–100x for large enough arrays.

The two habits this teaches:

1. Never write a for loop over batch dimensions in your inner training loop. If you find yourself iterating over batches inside a batch, your code is wrong. Reshape, broadcast, einsum.
2. np.einsum and torch.einsum are your friends. When you can't see how to broadcast, write the index notation explicitly and let the library figure it out. The string 'bnk,bkm->bnm' just names each axis with a letter: any letter that appears on the left but not after the -> is summed over (here k, exactly the inner dimension of a matmul), and a shared leading letter like b is carried along untouched — so this does one (n,k) @ (k,m) matmul for each of the b batch elements. You do not need to be fluent in this yet; it gets a proper treatment once batched tensors become routine. Reach for it only when broadcasting gets hard to read.

# Batched matmul without einsum: messy
A = np.random.randn(32, 10, 20)   # (batch, n, k)
B = np.random.randn(32, 20, 15)   # (batch, k, m)
C = A @ B                          # works, but you have to know matmul broadcasts the leading dims

# With einsum: explicit
C = np.einsum('bnk,bkm->bnm', A, B)

pandas is the library for reading CSVs, joining tables, filtering rows, and slicing columns. You use it before NumPy in every project, because real data lives in CSVs and TSVs and SQLite files, not in np.random.randn calls.

import pandas as pd

# Load
df = pd.read_csv("housing.csv")

# Inspect
df.head()             # first 5 rows
df.shape              # (n_rows, n_cols)
df.dtypes             # column types
df.describe()         # summary statistics per numeric column
df.info()             # null counts and types

# Select
df["price"]                          # one column, returns Series
df[["price", "area"]]                # multiple columns, returns DataFrame
df.iloc[0:5, :]                      # positional indexing
df.loc[df["price"] > 1e6, "area"]    # boolean mask + column select

# Mutate
df["price_per_sqft"] = df["price"] / df["area"]

# Aggregate
df.groupby("ocean_proximity")["price"].mean()

# Missing data
df.isna().sum()        # how many NaNs per column
df.dropna()            # drop rows with any NaN
df.fillna(df.median(numeric_only=True))   # fill with median

# To NumPy
X = df.drop(columns=["target"]).values
y = df["target"].values

That covers ≈80% of what you'll do with pandas in this curriculum. The other 20% is pd.merge (SQL-style joins), pd.pivot_table (Excel-style pivots), and pd.to_datetime (time series parsing).

Plot once. Look at the data. Adjust the model. Look at the data again. Matplotlib is the library for the looking-at part.

import matplotlib.pyplot as plt
import numpy as np

# Single plot
x = np.linspace(0, 10, 100)
y = np.sin(x)
plt.figure(figsize=(6, 4))
plt.plot(x, y, label="sin(x)")
plt.xlabel("x")
plt.ylabel("y")
plt.title("A sine wave")
plt.legend()
plt.grid(True)
plt.show()

# Multiple subplots
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
axes[0].scatter(np.random.randn(100), np.random.randn(100))
axes[0].set_title("Scatter")
axes[1].hist(np.random.randn(1000), bins=30)
axes[1].set_title("Histogram")
plt.tight_layout()
plt.show()

# Common ML plots
# Training curve
plt.plot(losses, label="train")
plt.plot(val_losses, label="val")
plt.xlabel("step")
plt.ylabel("loss")
plt.yscale("log")    # log scale for loss curves
plt.legend()

# Confusion matrix as a heatmap
import sklearn.metrics
cm = sklearn.metrics.confusion_matrix(y_true, y_pred)
plt.imshow(cm, cmap="Blues")
plt.colorbar()
plt.xlabel("predicted")
plt.ylabel("true")

For convenience, here's the boilerplate I paste into the first cell of every notebook in this curriculum.

# Common imports
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F
import sklearn

# Reproducibility
np.random.seed(42)
torch.manual_seed(42)

# Better plot defaults
plt.rcParams['figure.figsize'] = (8, 5)
plt.rcParams['figure.dpi'] = 100
plt.rcParams['axes.grid'] = True
plt.rcParams['grid.alpha'] = 0.3

# Device selection
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

Three habits this builds:

1. Seed RNGs when reproducibility matters. For learning and debugging, deterministic behavior helps you tell whether a code change caused a result. Production systems may also need reproducible tests and incident replay, while intentionally stochastic behavior should use an explicit, controlled source of randomness.
2. Always check your device. Putting tensors on the wrong device produces a runtime error. Putting them on CPU when you meant GPU produces a 50x slowdown silently.
3. Always set plot defaults once. Tiny plots are unreadable. The default 6x4 dpi 72 plot from matplotlib is too small. Bump it.