Implement common neural network layers in PyTorch

Implement the following neural-network building blocks from scratch in PyTorch (using only basic tensor ops such as matmul, sum, mean, var, exp, max, reshape, transpose, broadcasting). Do not call torch.nn.* or torch.nn.functional.* implementations of the same ops.

1. ReLU

• Input: tensor x of any shape.
• Output: y = max(x, 0).

2. Softmax (numerically stable)

• Input: tensor x and an integer dim.
• Output: softmax along dim.
• Must be stable for large/small values.

3. LayerNorm

• Input: x with shape [*, normalized_dim] (or more generally normalize over the last k dims), learnable gamma, beta, and eps.
• Output: normalized tensor.

4. BatchNorm (training and inference)

• Input: x (e.g., [N, C, H, W]), learnable gamma, beta, eps, momentum, plus running_mean and running_var.
• Implement:
  • Training forward: compute batch stats and update running stats.
  • Inference forward: use running stats.

5. RMSNorm

• Input: x, learnable gamma, eps.
• Normalize by RMS (no mean subtraction).

For each component, describe how you would verify correctness (e.g., compare to a reference implementation, run gradient checks / finite-difference checks, test edge cases such as fp16, tiny variance, and different tensor shapes).