Batch Normalization Forward Pass

Implement the forward pass of batch normalization using PyTorch tensors.

Batch normalization normalizes activations across the batch dimension:

$\hat{x}_i = \frac{x_i - \mu_B}{\sqrt{\sigma_B^2 + \epsilon}}$ $y_i = \gamma \hat{x}_i + \beta$

Where:

• $\mu_B = \frac{1}{m}\sum_{i=1}^{m} x_i$ (batch mean)
• $\sigma_B^2 = \frac{1}{m}\sum_{i=1}^{m} (x_i - \mu_B)^2$ (batch variance)
• $\gamma, \beta$ are learnable parameters
• $\epsilon$ is for numerical stability

Also track running mean/variance for inference mode.