Implement Neural Network Layers from Scratch in PyTorch

Your task is to implement several common neural network building blocks entirely from scratch in PyTorch. You are restricted to using only basic tensor operations such as matmul, sum, mean, var, exp, max, reshape, transpose, and broadcasting. Crucially, you must not use any implementations from torch.nn.* or torch.nn.functional.* for these operations.

Layers to Implement:

1. ReLU (Rectified Linear Unit)

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

2. Softmax (Numerically Stable)

   • Input: A tensor x and an integer dim specifying the dimension along which to apply Softmax.
   • Output: The Softmax applied along the specified dim.
   • Constraint: The implementation must be numerically stable for both very large and very small input values.

3. Layer Normalization (LayerNorm)

   • Input: A tensor x with shape [*, normalized_dim] (or more generally, normalize over the last k dimensions), along with learnable parameters gamma and beta, and a small constant eps for numerical stability.
   • Output: The normalized tensor.

4. Batch Normalization (BatchNorm) - Training and Inference Modes

   • Input: A tensor x (e.g., typically [N, C, H, W] for convolutional layers), learnable parameters gamma and beta, a small constant eps, a momentum value for updating running statistics, and stateful running_mean and running_var tensors.
   • Implementation Requirements:
     • Training Forward Pass: Compute the mean and variance of the current batch, use these to normalize the input, and update the running_mean and running_var using the provided momentum.
     • Inference Forward Pass: Use the stored running_mean and running_var (instead of batch statistics) to normalize the input.

5. Root Mean Square Normalization (RMSNorm)

   • Input: A tensor x, learnable parameter gamma, and a small constant eps.
   • Output: The normalized tensor.
   • Constraint: Normalize by the Root Mean Square (RMS) value, meaning no mean subtraction is performed.

Verification

For each implemented component, you must describe how you would verify its correctness. This could include strategies such as comparing against a reference implementation (e.g., torch.nn equivalents), performing gradient checks (e.g., finite-difference checks), and testing edge cases like fp16 precision, tiny variance values, and various tensor shapes and ranks.