Skip to content

← Performance step 1 of 7

Medium Primitives

Replace an elementwise loop

Jump to the task

A Python loop that applies a small operation to each tensor element repeats indexing and dispatch work. A whole-tensor expression can express the same calculation with fewer calls:

import torch

x = torch.tensor([1.0, 2.0, 3.0])
out = torch.empty_like(x)
for i in range(x.numel()):
    out[i] = x[i] * 2
print(out.tolist())    # [2.0, 4.0, 6.0]
print((x * 2).tolist())  # [2.0, 4.0, 6.0]

In the eager dispatch counter used by this course, the loop above takes 1 + 4*n operations for n elements, versus one multiplication dispatch for x * 2. For n=64, those counts are 257 and one. A dispatch is not necessarily a device kernel, and these counts are not a speedup measurement.

Look for independent work

Vectorization is useful when each element or row follows the same calculation. Reductions such as sum or max can combine a whole dimension, while broadcasting can apply each row's result back across that row.

A loop over large independent tensor operations may be perfectly reasonable. Some calculations also have sequential dependencies. Look for supported scans, reductions or batched operations where they fit, rather than assuming every loop must disappear. Measure the effect on the workload you care about.

Normalize rows with a denominator floor

For this exercise, first find each row's largest absolute value. Keep the reduced dimension so the resulting column of scales broadcasts across x. Use a minimum denominator of 1e-12 before dividing.

The floor avoids division by zero. It also changes the result for a nonzero row whose maximum magnitude is below 1e-12: that row will not be scaled up to maximum magnitude one. This is deliberate stabilized normalization.

Your task

def normalise_rows(x: torch.Tensor) -> torch.Tensor

The input is a nonempty two-dimensional float32 tensor with finite values and at least one column. For each row, divide by max(max(abs(row)), 1e-12). An all-zero row stays zero.

Use at most six dispatched operations, independent of the number of rows or columns. The harness returns values rounded to four decimals and a cheap flag for this dispatch budget. The amount of arithmetic still grows with the number of elements, even when the dispatch count stays fixed.

For rows [[1.0, 2.0], [0.0, 0.0]], return [[0.5, 1.0], [0.0, 0.0]]. A row [1e-14, -2e-14] becomes approximately [0.01, -0.02] because of the floor. The starter loops over rows and elements and omits the floor for tiny rows.