Skip to content
Medium Framework

Apply Along an Axis

Turn each row of positive values into proportions that sum to one. Divide each value by the total of its own row, without mixing values from other rows.

Your task

Implement row_normalize(x) for a floating-point matrix of shape (m, n), with n > 0 and strictly positive entries. Return the same shape, dtype and device, leaving x unchanged.

row_total[i] = sum(x[i, j] over all columns j)
output[i, j] = x[i, j] / row_total[i]

For [[1, 3, 6], [2, 2, 4]], the row totals are 10 and 8. The result is [[0.1, 0.3, 0.6], [0.25, 0.25, 0.5]]. Each row sums to one, up to floating-point rounding. A row with just one positive value becomes [1].

Keep a dimension for broadcasting

Sum across the columns, which is axis 1. Keep that axis with size 1 so that the totals have shape (m, 1). Dividing (m, n) by (m, 1) then applies each total across the corresponding row's columns.

Use x / x.sum(dim=1, keepdim=True) in PyTorch or x / x.sum(axis=1, keepdims=True) in JAX. Dropping the reduced axis gives shape (m,), which may fail to broadcast or divide along the wrong direction.

For these positive inputs, dividing by the row sum is L1 normalization because the sum of absolute values equals the sum of values. No zero-row fallback or added epsilon is needed under the input contract. This is not softmax: do not exponentiate the values before dividing.

The shape diagram shows how one total per row is reused across its columns.

Loading visualization…