Python and NumPy Basics for Deep Learning

The numerical examples in this series rely on arrays. This article introduces creation, dtype, shape, axes, and indexing, then connects them to a few neural-network operations. NumPy Vectorization and Broadcasting develops how array operations combine across shapes.

This article assumes basic Python and focuses on NumPy. The series uses functions with default and keyword arguments, dictionaries for named collections of parameters, f-strings for output, list comprehensions, enumerate and zip, and classes with __init__ once PyTorch modules appear. Type hints sometimes annotate intended inputs and outputs; Python does not enforce array shapes or dtypes from those hints alone. A lambda, used below, is a short way to define a function that returns one expression.

A working environment

A virtual environment gives the project its own package installation directory. Run the commands below in a terminal, not inside Python. The first two lines are for macOS or Linux shells such as bash and zsh. On Windows, create the environment with py -m venv .venv (or python -m venv .venv if that is how Python is installed), then use the activation command for your shell instead of source. After activation, the install command is the same.

python3 -m venv .venv
source .venv/bin/activate
# Windows PowerShell:  .\.venv\Scripts\Activate.ps1
# Windows cmd:         .venv\Scripts\activate.bat
python -m pip install numpy matplotlib jupyter

After installation, python -m notebook opens a notebook server using the active environment. You can also save the Python blocks in a .py file and run it with python filename.py. PyTorch installation is covered later because the appropriate build depends on the platform and hardware. If a compiled package fails to import after a NumPy upgrade, identify the failing package and follow its compatibility guidance; upgrading or reinstalling a compatible build may resolve it. A temporary numpy<2 pin only works with a compatible Python version—NumPy 1.26.4 supports Python 3.9–3.12. The NumPy troubleshooting guide explains these import errors.

The matrix shapes used throughout are introduced in Linear Algebra for Deep Learning: Only What You Actually Need.

Why not plain Python lists

A Python list stores references to Python objects, which may have different types. An ordinary numeric NumPy array stores fixed-size elements with a shared dtype in a data buffer, so many operations can process those elements in compiled code. A stride is the number of bytes to move along an axis to reach the next element. Views can share a buffer while using different strides, so slices and transposes need not be contiguous in memory.

import time
import numpy as np

n = 1_000_000
lst = list(range(n))
arr = np.arange(n, dtype=np.int64)

def best(f, reps=3):                      # both versions do the same work:
    out = []                              # double every element, then add them up
    for _ in range(reps):
        t = time.perf_counter(); f(); out.append(time.perf_counter() - t)
    return min(out)

assert sum(x * 2 for x in lst) == int((arr * 2).sum()) == n * (n - 1)

py  = best(lambda: sum(x * 2 for x in lst))
npt = best(lambda: (arr * 2).sum())
print(f"python {py*1000:.1f} ms | numpy {npt*1000:.1f} ms | {py/npt:.0f}x")
# python 60.7 ms | numpy 2.0 ms | 30x    (the ratio varies by machine and by run)

Both versions double the same integers and sum them; the assertion checks that their results agree. Array creation is outside the timed region, while the NumPy expression allocates a temporary for arr * 2. The reported times are the fastest of three trials, so the speed ratio is an illustration, not a portable benchmark. Array operations often reduce Python overhead for large batches. Loops remain useful for sequential work and for processing data in pieces that fit in memory.

Creating arrays

Run the blocks from here on in Python or a notebook, in the order they appear: later blocks reuse names defined earlier, such as np, A, and v. The first six lines below are bare expressions. In a REPL, running each one separately displays its value; a notebook normally displays only the final expression in a cell, and a script displays none of them. Wrap each expression in print() to see every result when running the whole block.

np.array([[1, 2], [3, 4]])        # from a nested list
np.zeros((2, 3))                  # all zeros, shape (2, 3)
np.ones((2, 3))                   # all ones
np.eye(3)                         # identity matrix
np.arange(0, 10, 2)               # [0 2 4 6 8]
np.linspace(0, 1, 5)              # [0. 0.25 0.5 0.75 1.]

rng = np.random.default_rng(42)   # seed fixes the initial generator state
print(np.round(rng.normal(size=(2, 3)), 3))
# [[ 0.305 -1.04   0.75 ]
#  [ 0.941 -1.951 -1.302]]

The first four calls construct arrays from values or prescribed patterns. arange(0, 10, 2) steps by 2 and excludes 10; linspace(0, 1, 5) requests five evenly spaced values and includes both endpoints by default. zeros, ones, and eye default to floating-point values.

np.random.default_rng(seed) creates an explicit generator instead of changing a shared global random state. Recreating it with the same seed and making the same calls in the same environment reproduces the sequence; another call on the existing generator advances that sequence. A seed alone does not guarantee identical results across NumPy versions or different execution environments.

dtype: storage, precision, and overflow

Every array has one element type. Assigning an in-range floating-point value into an integer array truncates toward zero. Fixed-width integer array arithmetic can also wrap around on overflow, as the int8 example below shows. Division is a separate matter: / on an integer array returns floats, // floors, and an in-place a /= 2 raises rather than truncating.

a = np.array([1, 2, 3], dtype=np.int64) # explicit 64-bit integers
b = np.array([1.0, 2.0, 3.0])        # inferred float64
print(a.dtype, b.dtype)              # int64 float64

x = np.array([5, 10])
x[0] = 2.7                           # silently truncated
print(x)                             # [ 2 10]

small = np.array([100], dtype=np.int8)
print(small + np.int8(100))          # [-56], wraps around

print(np.array([1, 2, 3], dtype=np.float32).dtype)   # float32
print(b.astype(np.float32).dtype)                    # float32

Without an explicit dtype, NumPy infers a type from the values, and the default integer width can depend on the platform and NumPy version. The example specifies int64 so the printed type is unambiguous. For continuous model inputs and weights, use floating point, and start from float32, which uses half the storage of float64 and is a common default for training. Move to higher precision when the numerical requirements justify it — gradient checking is the clearest case in this series, but it is not the only one. Class-index targets are integers, but other targets can be floating point; the exact target dtype depends on the loss function, and for PyTorch’s CrossEntropyLoss class indices are int64 while probability targets are floating point.

Shape and axes

An array’s shape lists the lengths of its axes, ndim counts those axes, and size counts all elements. For a (2, 3) matrix, axis 0 runs through the two rows and axis 1 through the three columns. A reduction, such as a sum or mean, combines values along selected axes.

A = np.array([[1., 2., 3.],
              [4., 5., 6.]])        # shape (2, 3)

print(A.shape, A.ndim, A.size)      # (2, 3) 2 6
print(A.sum())                      # 21.0        everything
print(A.sum(axis=0))                # [5. 7. 9.]  collapse rows -> (3,)
print(A.sum(axis=1))                # [ 6. 15.]   collapse cols -> (2,)
print(A.sum(axis=1, keepdims=True)) # [[ 6.] [15.]]  keeps 2-D: (2, 1)

For the matrix above, the first column adds to 1 + 4 = 5, and the first row to 1 + 2 + 3 = 6. Read axis=0 as “move down the rows and combine them,” so the row axis disappears. keepdims=True leaves a length-1 axis behind, and later articles use it constantly to keep results aligned for the next operation.

Reshaping changes the shape without changing the element count. It returns a view when the existing strides allow one and copies when they do not; flatten() below always copies. The -1 placeholder means “infer this one.”

v = np.arange(6)
print(v.shape, v.reshape(2, 3).shape, v.reshape(-1, 2).shape)  # (6,) (2, 3) (3, 2)
print(A.T.shape)                                               # (3, 2)  transpose
print(v.reshape(2, 3).flatten().shape)                         # (6,)

col = v.reshape(-1, 1)      # (6, 1) column
row = v.reshape(1, -1)      # (1, 6) row
print(col.shape, row.shape)

(6,) has one axis and is neither a two-dimensional row nor a column. Its transpose still has shape (6,). The reshaped arrays (6, 1) and (1, 6) place the length-six axis differently; for example, six elements reshaped with (-1, 2) need three rows because 6 / 2 = 3. These distinctions determine which axes can align in later array operations.

Indexing and slicing

Two pieces of syntax appear below. A slice start:stop includes start and excludes stop, so 1:3 selects positions 1 and 2; a negative index counts from the end, so -1 is the last one. And % is the remainder operator, so M % 2 == 0 is a boolean array marking the even entries.

M = np.arange(12).reshape(3, 4)
# [[ 0  1  2  3]
#  [ 4  5  6  7]
#  [ 8  9 10 11]]

print(M[1, 2])          # 6        single element
print(M[0])             # [0 1 2 3]        first row
print(M[:, 1])          # [1 5 9]          second column -> 1-D!
print(M[:, 1:2])        # [[1] [5] [9]]    same column, kept 2-D
print(M[0:2, 1:3])      # [[1 2] [5 6]]    submatrix
print(M[-1])            # [ 8  9 10 11]    last row

mask = M % 2 == 0
print(M[mask])          # [ 0  2  4  6  8 10]   boolean selection
print(M[[0, 2]])        # rows 0 and 2         fancy indexing

Compare M[:, 1] with M[:, 1:2]: the integer index drops that axis, while the slice keeps it. A bare : selects the full axis. Here M[mask] returns a one-dimensional array of the selected values in row order. keepdims is not an indexing option — it belongs to reduction operations like sum — but it serves the same purpose there, leaving a length-1 axis in place instead of removing it.

Views versus copies

A basic slice is a view into the original memory, so writing through it changes the original. Reading with boolean or fancy indexing produces a copy. Assigning into a fancy index is different again: orig[[1, 2, 3]] = 99 writes to the original, while orig[[1, 2, 3]][0] = 99 modifies a temporary copy and throws it away. This distinction matters when modifying a subset of training data: decide whether the original array should change.

orig = np.arange(5)
view = orig[1:4]
view[0] = 99
print(orig)             # [ 0 99  2  3  4]   original changed

orig = np.arange(5)
copy = orig[[1, 2, 3]]  # fancy indexing copies
copy[0] = 99
print(orig)             # [0 1 2 3 4]        original untouched

safe = orig[1:4].copy() # be explicit when you mean a copy

Elementwise operations and useful functions

x = np.array([-2., -0.5, 0., 1., 3.])

print(np.exp(x).round(3))        # [0.135 0.607 1.    2.718 20.086]
print(np.maximum(x, 0))          # [0.  0.  0.  1.  3. ]   this is ReLU
print(np.where(x > 0, 1., 0.))   # [0. 0. 0. 1. 1. ]       its local gradient mask
print(np.abs(x).sum())           # 6.5
print(x.argmax(), x.max())       # 4 3.0
print(np.clip(x, -1, 1))         # [-1.  -0.5  0.   1.   1. ]

Two lines implement pieces used in neural networks: np.maximum(x, 0) is the ReLU activation, and np.where(x > 0, 1., 0.) produces its local gradient mask. ReLU is not differentiable at zero, and the input array above contains 0.0 — the np.where line assigns 0 there by convention rather than computing a derivative. The mask is not itself the gradient of the input: backpropagation multiplies the incoming gradient by it elementwise, dx = upstream_gradient * np.where(x > 0, 1.0, 0.0).

Python features this series relies on

The next example uses a dictionary to name parameter arrays, then pairs consecutive layer sizes. The zero-filled weights illustrate storage and shape only; initializing every weight of a hidden layer to zero generally prevents its neurons from learning different features, so the training examples use other initializations.

params = {"W1": np.zeros((4, 3)), "b1": np.zeros((4, 1))}
for name, value in params.items():
    print(f"{name}: {value.shape}")
# W1: (4, 3)
# b1: (4, 1)

layer_dims = [784, 256, 10]
for l, (n_in, n_out) in enumerate(zip(layer_dims[:-1], layer_dims[1:]), 1):
    print(f"layer {l}: {n_in} -> {n_out}")
# layer 1: 784 -> 256
# layer 2: 256 -> 10

layer_dims[:-1] omits the last size and layer_dims[1:] omits the first. zip pairs them: (784, 256) and (256, 10). enumerate(..., 1) numbers those pairs starting at 1. Each pair gives the input and output sizes for one layer. The from-scratch articles use this dictionary-of-arrays pattern to store network parameters and extend it to networks with more layers.

The first network built from these pieces is in Neural Networks for Beginners: From Architecture to Forward Propagation.

Inspecting an array before using it

Check the shapes you expect as you build each operation. An array with a shape you did not expect is a frequent failure in the articles that follow. Some mismatches raise errors; others broadcast successfully while computing something you did not intend. The helper below summarizes nonempty real numeric arrays. Shape and range checks are useful diagnostics, but they do not prove that a calculation is correct.

def describe(name, a):
    print(f"{name:8s} shape={str(a.shape):12s} dtype={a.dtype} "
          f"min={a.min():.3f} max={a.max():.3f}")

describe("X", np.random.default_rng(0).normal(size=(3, 100)))
# X        shape=(3, 100)     dtype=float64 min=-3.106 max=3.066

You now have the operations needed to create parameter arrays, select batches, and inspect intermediate results in the forward-propagation example linked above.

The training loop this fits into is built in Building a PyTorch Training Loop with nn.Module, Dataset, and DataLoader.

Exercises

1. Axis reasoning. For A of shape \((2,3,4)\), state the shape of A.sum(axis=0), A.sum(axis=1), A.sum(axis=(0,2)), and A.sum(axis=1, keepdims=True) before running anything.

You should get: three shapes with one or two axes removed, and one that keeps three axes.

import numpy as np
A = np.zeros((2,3,4))
for s in [A.sum(axis=0), A.sum(axis=1), A.sum(axis=(0,2)), A.sum(axis=1, keepdims=True)]:
    print(s.shape)
# (3, 4)
# (2, 4)
# (3,)
# (2, 1, 4)

2. The view trap. Predict the printed output, then run it. Explain in one sentence why the two cases differ.

You should get: one array modified and one untouched. Slicing and fancy indexing behave differently.

import numpy as np
a = np.arange(6); a[2:5][0] = 99; print(a)
b = np.arange(6); b[[2,3,4]][0] = 99; print(b)
# [ 0  1 99  3  4  5]
# [0 1 2 3 4 5]

3. Implement ReLU and its local gradient rule. Write both in one line each using NumPy, then compare against a numerical derivative at \(z=-1,0,2\). Explain why the two disagree at exactly \(z=0\) and which convention your code is using.

You should get: agreement at \(-1\) and \(2\), and a discrepancy at \(0\).

import numpy as np
relu  = lambda z: np.maximum(z, 0)
drelu = lambda z: np.where(np.asarray(z) > 0, 1.0, 0.0)
h = 1e-6
for z in (-1., 0., 2.):
    num = (relu(z+h) - relu(z-h)) / (2*h)
    print(z, drelu(z), round(float(num), 6))
# -1.0 0.0 0.0
# 0.0 0.0 0.5
# 2.0 1.0 1.0