Per-Channel Batch Normalization

Implement batchnorm2d(x, gamma, beta, eps=1e-5) that performs training-mode batch normalization for a 4-D input.

Shapes:

• x has shape (N, C, H, W).
• gamma (scale) and beta (shift) each have shape (C,).

For each channel $c$ independently, compute the mean and variance over the batch, height, and width dimensions, normalize, then scale and shift:

$\hat{x}_{n c h w} = \frac{x_{n c h w} - \mu_c}{\sqrt{\sigma_c^2 + \epsilon}}, \quad y_{n c h w} = \gamma_c \, \hat{x}_{n c h w} + \beta_c.$

You do not need to track running statistics here - we only care about the training-mode forward pass.

Hints

• The mean and variance reduce over three axes simultaneously: the batch axis and the two spatial axes.
• Use keepdim=True when reducing so broadcasting against the original tensor works without manual reshapes.
• gamma and beta are 1-D; they need to be reshaped so they broadcast against (N, C, H, W).
• For the variance, use the population (biased) formula, not the sample (unbiased) one.