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).