RMS Normalization

Implement RMS (Root Mean Square) Normalization for a 2D tensor.

Normalize the input by dividing each element by the root mean square of the features in each sample. More formally, compute:

$$\text{y} = \frac{x}{\sqrt{\text{mean}(x^2) + \epsilon}}$$

where the mean is computed along the feature dimension for each sample in the batch independently. $\epsilon$ is a small value added to the denominator for numerical stability.

Input:

• Tensor $\text{X}$ of shape $(\text{B}, \text{N})$ is the input data where $\text{B}$ = batch size and $\text{N}$ = number of features

Output:

• Tensor $\text{Y}$ with the same shape as input (normalized data)

Notes:

• For each sample, the RMS is calculated over the feature dimension (dimension 1).
• Use $\epsilon = 10^{-5}$ for numerical stability.

Test Case Sizes

• shape=(1024, 1024)
• shape=(1024, 4096)
• shape=(2048, 8192)
• shape=(512, 16384)