NumPy basics for ML in Practical AI & Machine Learning

NumPy basics for ML in Practical AI & Machine Learning. Arrays are the foundation of every ML library.

The ndarray: Fast, typed, multi-dimensional

Every machine learning library we've used in production—scikit-learn, TensorFlow, PyTorch—is built on top of NumPy arrays. The ndarray is the core data structure, and understanding it deeply will save you hours debugging shape mismatches and mysterious performance bottlenecks.

An ndarray is fundamentally different from a Python list. It's a contiguous block of memory containing elements of a single data type. This homogeneity is what makes NumPy fast: operations can be compiled to machine code without type-checking overhead. When you create an array, NumPy allocates memory once and keeps it organized. A Python list, by contrast, stores pointers to objects scattered across memory.

Here's what that looks like in practice:

import numpy as np

# Create arrays with explicit dtype
x = np.array([[1, 2], [3, 4]], dtype=np.float32)
print(x.dtype)        # float32
print(x.shape)        # (2, 2)
print(x.nbytes)       # 16 bytes (4 elements × 4 bytes)

# Compare with a list
py_list = [[1, 2], [3, 4]]
# In memory: scattered objects, references, overhead

# Create larger arrays efficiently
zeros = np.zeros((1000, 1000), dtype=np.float64)
ones = np.ones((100, 50), dtype=np.int32)
range_array = np.arange(0, 100, 5)  # 0, 5, 10, ..., 95

Notice we always specify dtype. In production ML pipelines, we typically use float32 for model inputs (GPU-friendly, smaller memory) and float64 when we need precision in statistical calculations. Getting the dtype wrong can silently corrupt your model's output or waste GPU memory.

Vectorize instead of looping

The cardinal rule of NumPy performance is simple: avoid Python loops. Every loop iteration has Python interpreter overhead. Vectorized operations—applied to entire arrays at once—run compiled C code underneath and are orders of magnitude faster.

Consider a real scenario: normalizing a batch of feature vectors. You have 10,000 samples, each with 784 features (like MNIST images flattened).

import numpy as np
import time

# Sample data: 10,000 samples × 784 features
X = np.random.randn(10000, 784)

# ❌ SLOW: Python loop
def normalize_slow(X):
    result = []
    for row in X:
        normalized = row / np.linalg.norm(row)
        result.append(normalized)
    return np.array(result)

start = time.time()
X_slow = normalize_slow(X)
print(f"Loop version: {time.time() - start:.3f}s")

# ✅ FAST: Vectorized
def normalize_fast(X):
    norms = np.linalg.norm(X, axis=1, keepdims=True)
    return X / norms

start = time.time()
X_fast = normalize_fast(X)
print(f"Vectorized version: {time.time() - start:.3f}s")

The vectorized version is typically 50–100× faster. The difference between a 10-second preprocessing step and a 100-millisecond one matters in development iteration speed.

When you feel the urge to write a loop, ask yourself: "Is there a NumPy function that does this?" Usually there is. Use axis parameters to specify which dimension to operate along, and use keepdims=True to preserve dimensions for broadcasting.

Broadcasting applies operations across shapes

Broadcasting is NumPy's mechanism for applying operations to arrays of different shapes. It's powerful but also a common source of bugs if you don't understand it deeply.

The broadcasting rule is straightforward: NumPy aligns arrays from the right and expands dimensions where one array has size 1 (or a dimension is missing). Let's walk through a typical ML scenario—mean-centering a batch of data:

import numpy as np

# Batch of 32 samples, 10 features each
X = np.random.randn(32, 10)

# Compute the mean of each feature (across samples)
feature_means = X.mean(axis=0)  # Shape: (10,)
print(feature_means.shape)  # (10,)

# Subtract means from each sample
# X is (32, 10), feature_means is (10,)
# Broadcasting aligns from the right:
#   (32, 10)
#   (   10,)  ← treated as (1, 10), then expanded to (32, 10)
X_centered = X - feature_means

print(X_centered.shape)  # (32, 10)
print(np.allclose(X_centered.mean(axis=0), 0))  # True (nearly zero)

This is clean and efficient. But broadcasting can surprise you if you're not careful:

import numpy as np

# ⚠️ Subtle broadcasting error
a = np.array([[1, 2, 3]])  # Shape: (1, 3)
b = np.array([[1], [2], [3]])  # Shape: (3, 1)

# Broadcasting expands both to (3, 3)
result = a + b
print(result.shape)  # (3, 3)
print(result)
# [[2, 3, 4],
#  [3, 4, 5],
#  [4, 5, 6]]

If your matrix dimensions don't match your expectations, you'll get the wrong answer silently. This is why the next habit is critical.

Check shape before every operation

This is our primary defense against silent errors. Make checking .shape a reflex, especially when reading or transforming data.

In a typical ML pipeline, we load data, split it, preprocess it, and feed it to a model. At each step, print the shape:

import numpy as np

# Load raw data
X_raw = np.load('features.npy')
print(f"Raw shape: {X_raw.shape}")  # (10000, 784)

# Train-test split
train_size = int(0.8 * len(X_raw))
X_train = X_raw[:train_size]
X_test = X_raw[train_size:]
print(f"Train shape: {X_train.shape}, Test shape: {X_test.shape}")
# Train shape: (8000, 784), Test shape: (2000, 784)

# Standardize (per-feature mean and std)
mean = X_train.mean(axis=0, keepdims=True)  # (1, 784)
std = X_train.std(axis=0, keepdims=True)    # (1, 784)
X_train_std = (X_train - mean) / std
X_test_std = (X_test - mean) / std  # Use train stats on test

print(f"Standardized train: {X_train_std.shape}")
print(f"Standardized test: {X_test_std.shape}")

# Flatten for a dense layer
X_train_flat = X_train_std.reshape(X_train_std.shape[0], -1)
print(f"Flattened train: {X_train_flat.shape}")  # (8000, 784)

Use keepdims=True liberally when computing statistics; it makes broadcasting explicit and prevents shape surprises. Use -1 in reshape to let NumPy infer that dimension.

Common mistakes

Forgetting axis when reducing: If you call X.mean() instead of X.mean(axis=0) or X.mean(axis=1), you'll reduce the entire array to a scalar. On a (1000, 50) matrix, this is a silent catastrophe if you expected (1000,) or (50,).

Modifying array copies unintentionally: Slicing an array (like X[:100]) returns a view, not a copy. If you modify it, you modify the original. Use .copy() when you need independence.

Dtype coercion in operations: If you divide an integer array by a float, NumPy promotes to float, which can silently consume extra memory. Be explicit: cast to the dtype you want upfront.

Off-by-one errors with arange vs linspace: np.arange(0, 10, 2) gives [0, 2, 4, 6, 8] (excludes 10). np.linspace(0, 10, 6) gives 6 evenly spaced points including both endpoints. For range-based operations, arange is usually right; for sampling, linspace is safer.

Key takeaways

• ndarray is a contiguous, typed, multi-dimensional array—the foundation of all ML libraries. Always specify dtype explicitly.
• Vectorization (operations on whole arrays) is 50–100× faster than Python loops. Use axis parameters and keepdims=True to control dimensions.
• Broadcasting automatically expands dimensions during operations. Understand the rules: align from the right, expand size-1 dimensions.
• Check .shape at every transformation step. Silent shape mismatches are the most insidious bugs in ML pipelines.
• Learn the core functions: mean, std, sum, dot, reshape, concatenate, argmax. These cover 80% of preprocessing work.