Skip to content

← Performance step 7 of 7

Hard Primitives

Vectorize a row-scoring pipeline

Jump to the task

This exercise combines row reductions, broadcasting and comparisons. The starter computes each row separately, extracts Python numbers, and rebuilds a tensor from the resulting list. A vectorized version can keep the calculation in tensors throughout.

Follow the shapes

For x with shape (rows, columns), reducing absolute values with sum(dim=1, keepdim=True) gives a scale tensor of shape (rows, 1). It can broadcast back across x. A weights vector of shape (columns,) then broadcasts across rows. Summing the weighted values along dimension one gives one total per row.

Comparing all those totals with a threshold produces a boolean tensor. Converting that tensor with float() gives float32 zeros and ones without extracting a Python scalar for each row.

import torch

x = torch.tensor([[1.0, 2.0], [-3.0, 1.0]])
weights = torch.tensor([1.0, 1.0])
scale = x.abs().sum(dim=1, keepdim=True).clamp(min=1e-12)
totals = (x / scale * weights).sum(dim=1)
flags = (totals > 0.5).float()
print(flags.tolist())  # [1.0, 0.0]

Keep the denominator rule explicit

Divide by the row's absolute sum floored at 1e-12. Zero rows remain zero, and nonzero rows below that floor are only partially normalized. The starter branches only on zero, so it differs for those tiny rows as well as exceeding the operation budget. The denominator floor is part of the requested result.

Avoid per-row scalar extraction

The starter calls float(scale) to make a Python decision and total.item() to build its list. With CUDA data, reading those values can make the host wait for their computation. Rebuilding a tensor from a list also uses the current default device rather than automatically following x.

A tensor comparison needs neither a Python branch nor per-row extraction. This does not imply that arbitrary tensor code never synchronizes. The exercise checks one specific scalar-extraction dispatch, not every possible transfer or wait.

Your task

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

Inputs are finite float32 tensors on the same device. x is a nonempty matrix with at least one column, and weights has one entry per column. For each row:

  1. Divide by max(sum(abs(row)), 1e-12).
  2. Multiply by weights and sum the products.
  3. Return 1.0 if that sum is strictly greater than threshold, otherwise 0.0.

Return a float32 vector on x's device and leave x unchanged. Use at most eight dispatched operations and no scalar-extraction dispatches inside the function. The harness runs on CPU, then checks the flags, operation budget and unchanged input. These checks do not measure GPU runtime or all forms of synchronization.

For the example above, return [1.0, 0.0]. Equality with the threshold must produce zero. Reduce and broadcast over whole rows rather than rebuilding the output through a Python list.