Compute the population mean and variance of a 1D float tensor in a single pass over the data. For a tensor x of length N, write the scalar mean and variance into mean_out[0] and var_out[0]:
μ=N1i=0∑N−1x[i],σ2=N1i=0∑N−1x[i]2−μ2This problem teaches two-statistic fusion: each program accumulates ∑x and ∑x2 for its tile in the same reduction pass, then atomic-adds both partials into scratch buffers. The host finalizes by dividing and applying the E[x2]−E[x]2 identity.
Input: x = [1.0, 2.0, 3.0, 4.0]
Output: mean = [2.5], var = [1.25]
Input: x = [0.0, 0.0, 0.0, 0.0]
Output: mean = [0.0], var = [0.0]
Use the identity σ2=E[x2]−μ2 so one pass over x produces both statistics. Reduce tl.sum(x) and tl.sum(x * x) in the same tile.
Allocate two scalar scratch tensors with torch.zeros(1, device=’cuda’) inside solve. Pass their pointers to the kernel and atomic-add into each.
Finalize on the host after the launch: divide both sums by N, subtract μ2, and copy into the output tensors. Do not call .item().
torch.float32 on cuda and contiguous; outputs each have shape (1,).atol=1e-2, rtol=1e-2 for both outputs to absorb float32 sum-of-squares drift at large N.Accepts: array
Submit a correct solution to see your kernel performance.
Compute the population mean and variance of a 1D float tensor in a single pass over the data. For a tensor x of length N, write the scalar mean and variance into mean_out[0] and var_out[0]:
μ=N1i=0∑N−1x[i],σ2=N1i=0∑N−1x[i]2−μ2This problem teaches two-statistic fusion: each program accumulates ∑x and ∑x2 for its tile in the same reduction pass, then atomic-adds both partials into scratch buffers. The host finalizes by dividing and applying the E[x2]−E[x]2 identity.
Input: x = [1.0, 2.0, 3.0, 4.0]
Output: mean = [2.5], var = [1.25]
Input: x = [0.0, 0.0, 0.0, 0.0]
Output: mean = [0.0], var = [0.0]
Use the identity σ2=E[x2]−μ2 so one pass over x produces both statistics. Reduce tl.sum(x) and tl.sum(x * x) in the same tile.
Allocate two scalar scratch tensors with torch.zeros(1, device=’cuda’) inside solve. Pass their pointers to the kernel and atomic-add into each.
Finalize on the host after the launch: divide both sums by N, subtract μ2, and copy into the output tensors. Do not call .item().
torch.float32 on cuda and contiguous; outputs each have shape (1,).atol=1e-2, rtol=1e-2 for both outputs to absorb float32 sum-of-squares drift at large N.Accepts: array
Submit a correct solution to see your kernel performance.