asc2.rms_norm

asc2.rms_norm(input: LocalTensor, gamma: LocalTensor, epsilon: PlainValue | float) → LocalTensor

Computes Root Mean Square Layer Normalization of input.

RMSNorm normalizes the input by the root mean square and scales by learnable parameters gamma. This is commonly used in transformer architectures as an alternative to LayerNorm.

The supported data types for the inputs are: float16, float32.

Parameters:

  • input – The input tensor to normalize (1D or 2D)

  • gamma – The scale parameter tensor (1D, same length as last dimension of input)

  • epsilon – Small constant added for numerical stability

Returns:

The normalized tensor with same shape as input

Return type:

LocalTensor

Raises:

  • TypeError – If input or gamma is not a LocalTensor

  • RuntimeError – If input dtype is not supported, input has more than 2 dimensions, or gamma dtype is not supported

Examples

Apply RMSNorm to a 2D tensor:

input = asc2.copy_in(x_gm, [0, 0], [32, 128])
gamma = asc2.copy_in(gamma_gm, [0], [128])
output = asc2.rms_norm(input, gamma, 1e-5)

Apply RMSNorm to a 1D tensor:

input = asc2.copy_in(x_gm, [0], [128])
gamma = asc2.copy_in(gamma_gm, [0], [128])
output = asc2.rms_norm(input, gamma, 1e-6)