Skip to content

← Performance step 7 of 7

Hard Primitives

Everything at once

The last problem in the course. The starter is code you will genuinely find, and it is wrong in five ways you have now met individually.

def score(x, weights, threshold):
    out = []
    for i in range(x.size(0)):
        row = x[i]
        scale = torch.tensor(0.0)
        for j in range(row.numel()):
            scale = scale + abs(row[j])
        normed = row / scale if float(scale) > 0 else row
        weighted = normed * weights
        total = weighted.sum()
        out.append(total.item())
    result = torch.tensor(out)
    return (result > threshold).float()

Correct. Roughly a thousand dispatched operations on a small input, and it scales with the data rather than with the shape.

What is wrong, and which track said so

  1. A Python loop over rows, and another over elements. Track 12: the inner loop dispatches a kernel per element and the arithmetic was never the cost.
  2. A per-row reduction written by hand. Track 4: abs().sum(dim=1) does every row at once, and keepdim=True gives the shape to divide by.
  3. .item() inside the loop. Track 9: a host synchronisation per row, which serialises the whole thing on a GPU.
  4. A rebuilt Python list, then torch.tensor(out). Track 11: joining allocates, and here it also drags every value back to the host and forward again.
  5. float(scale) > 0 as a branch. Track 12 and track 4: a Python-level test per row, replaceable by clamp on the denominator with no branch at all.

Each of those is a track in this course. Fixing them together is what the course was for.

Your task

def score(x: torch.Tensor, weights: torch.Tensor, threshold: float) -> torch.Tensor

Same result. Row i of x is divided by the sum of its absolute values (leaving an all-zero row alone), multiplied elementwise by weights, summed, and compared against threshold. Return a float tensor of ones and zeros.

Constraints:

  • at most eight dispatched operations, whatever the input size
  • no host synchronisations
  • x unchanged

There is a version that satisfies all three and is four lines.