def call(self, x):
        """Applies RMS normalization to the input tensor.

        Args:
            x: Input tensor of shape (batch_size, input_dim).

        Returns:
            The RMS-normalized tensor of the same shape (batch_size, input_dim),
            scaled by the learned `scale` parameter.
        """
        return ops.rms_normalization(
            x, scale=self.scale, axis=self.axis, epsilon=self.epsilon
        )
Applies RMS normalization to the input tensor.

        Args:
            x: Input tensor of shape (batch_size, input_dim).

        Returns:
            The RMS-normalized tensor of the same shape (batch_size, input_dim),
            scaled by the learned `scale` parameter.
        
keras/src/layers/normalization/rms_normalization.py