PyTorch: broadcasting and reduction axes

pythontorchindependent-practice

Saved on this device

Full lesson · 4 learning steps Open any exercise for focused practice

Step 1 of 4 · intro card

PyTorch: broadcasting and reduction axes

Concept and contract

Center feature columns without accidentally mixing examples and features.

Broadcasting aligns compatible dimensions from the right and logically expands size-one axes without requiring a materialized copy for every repeated value.

That convenience can also hide a semantic mistake. For a matrix shaped examples by features, a reduction over dimension zero computes one value per feature, while dimension one computes a different quantity per example.

Mechanism and worked example

The helper computes feature means with dim=0 and keepdim=True, retaining shape one by features.

The helper computes feature means with dim=0 and keepdim=True, retaining shape one by features. Subtracting that result from the batch uses an explicit broadcast.

The returned means and centered matrix let the learner inspect both the reduction and its consequence. Nonempty two-dimensional input is required so a missing axis or empty reduction does not quietly propagate meaningless values.

Complete worked example

python

from bootstrap_lab import torch_setup
torch = torch_setup()

def centered(rows):
    x = torch.tensor(rows, dtype=torch.float32)
    if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] == 0:
        raise ValueError('nonempty matrix required')
    means = x.mean(dim=0, keepdim=True)
    result = x - means
    return {'means': means.tolist(), 'centered': result.tolist()}

Reasoning: The helper computes feature means with dim=0 and keepdim=True, retaining shape one by features. Subtracting that result from the batch uses an explicit broadcast.

Counterexamples and common mistakes

The counterexample reduces over dimension one, producing row-centered data instead of feature-centered data..

The counterexample reduces over dimension one, producing row-centered data instead of feature-centered data.

Some square or uniform examples can make the mistake look plausible, so the fixtures use asymmetric values. Do not fix a shape error by adding random unsqueeze calls until subtraction runs; first identify which axis represents examples and which represents measured features.

Counterexample: a runnable framework mistake

python

from bootstrap_lab import torch_setup
torch = torch_setup()

def centered(rows):
    x = torch.tensor(rows, dtype=torch.float32)
    if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] == 0:
        raise ValueError('nonempty matrix required')
    means = x.mean(dim=1, keepdim=True)
    result = x - means
    return {'means': means.tolist(), 'centered': result.tolist()}

Reasoning: The counterexample reduces over dimension one, producing row-centered data instead of feature-centered data.

Independent practice: one bounded framework change

After centering, divide each feature by its population standard deviation computed with correction=0.

After centering, divide each feature by its population standard deviation computed with correction=0. Replace a zero standard deviation with one so constant features remain zero instead of producing NaN.

Preserve the original means field. This is a small numerical policy exercise, not permission to fit normalization on a future validation set; train-only statistics appear in a later lesson.

Step 1: Scale centered feature columns with safe population standard deviations. Step 2: Compute std on the same feature axis.

Practice starter: solve before reading the solution

python

from bootstrap_lab import torch_setup
torch = torch_setup()

def centered(rows):
    # Remove this placeholder, then adapt the retained baseline:
    # Scale centered feature columns with safe population standard deviations.
    raise NotImplementedError
    x = torch.tensor(rows, dtype=torch.float32)
    if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] == 0:
        raise ValueError('nonempty matrix required')
    means = x.mean(dim=0, keepdim=True)
    result = x - means
    return {'means': means.tolist(), 'centered': result.tolist()}

Reasoning: Scale centered feature columns with safe population standard deviations. Compute std on the same feature axis.

Hint 1: Use correction=0 and replace zero scale with one.
Hint 2: Keep finite output for constant columns.

Solution, tests and trade-offs

Baseline fixtures include two nonuniform columns, a single example and invalid one-dimensional input.

Baseline fixtures include two nonuniform columns, a single example and invalid one-dimensional input. Practice adds a constant feature and compares standardized values using finite tolerances.

Check both shape and values: broadcasting can produce a valid result with the wrong semantics. The fixture uses tiny deterministic arrays and does not claim statistical suitability for a real dataset.

Practice solution and acceptance checks

python

from bootstrap_lab import torch_setup
torch = torch_setup()

def centered(rows):
    x = torch.tensor(rows, dtype=torch.float32)
    if x.ndim != 2 or x.shape[0] == 0 or x.shape[1] == 0:
        raise ValueError('nonempty matrix required')
    means = x.mean(dim=0, keepdim=True)
    scale = x.std(dim=0, correction=0, keepdim=True)
    scale = torch.where(scale == 0, torch.ones_like(scale), scale)
    result = (x - means) / scale
    return {'means': means.tolist(), 'centered': result.tolist()}

Reasoning: Baseline fixtures include two nonuniform columns, a single example and invalid one-dimensional input. Practice adds a constant feature and compares standardized values using finite tolerances.

Interview reasoning and next step

When features have different units or distributions, choose normalization deliberately and fit its parameters only on the appropriate training population.

When features have different units or distributions, choose normalization deliberately and fit its parameters only on the appropriate training population. Save those parameters alongside model state for inference.

Broadcasting remains useful, but naming axes and checking invariants should come first. A column-centered training matrix should have approximately zero column means, while that need not hold for unseen validation examples.

Optional stretch — after the core task works: Save train-only normalization parameters for a separate validation batch.

Computer practice: https://applydjinn.com/typorithm/python-frameworks-lab-20261105.zip. Install the pinned dependencies in an isolated environment, run verify.py --track torch, then edit the retained exercise baseline and check it with --practice pt-broadcast. The lab uses owned in-memory databases, local framework clients and tiny CPU tensors. It never reads the working ApplyDjinn database, contacts production feeds or downloads models.

Center feature columns without accidentally mixing examples and features.