asc.experimental.asctile.layer_norm

asc.experimental.asctile.layer_norm(input: LocalTensor, gamma: LocalTensor, beta: LocalTensor, epsilon: PlainValue | float, output_rstd: bool = True) → Tuple[LocalTensor, LocalTensor, LocalTensor]

Computes Layer Normalization of input.

LayerNorm normalizes the input by subtracting the mean and dividing by the standard deviation, then scales by learnable parameter gamma and shifts by learnable parameter beta. This is commonly used in transformer architectures.

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

Parameters:

  • input – The input tensor to normalize with shape [A, R] for 2D or [R] for 1D

  • gamma – The scale parameter tensor with shape [R] (same length as last dimension of input)

  • beta – The shift parameter tensor with shape [R] (same length as last dimension of input)

  • epsilon – Small constant added for numerical stability

  • output_rstd – If True (default), the third output is the reciprocal of standard deviation (rstd). If False, the third output is the variance. Note: on C220 architecture, this parameter is ignored and rstd is always returned.

Returns:

A tuple containing:

  • output: The normalized tensor with same shape and dtype as input

  • mean: The mean tensor with shape [A] for 2D input or [1] for 1D input (always float32)

  • rstd or variance: If output_rstd=True, returns reciprocal of standard deviation; otherwise returns variance. Shape [A] for 2D input or [1] for 1D input (always float32)

Return type:

Tuple[LocalTensor, LocalTensor, LocalTensor]

Raises:

  • TypeError – If input, gamma, or beta is not a LocalTensor

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

Examples

Apply LayerNorm to a 2D tensor:

input = asctile.copy_in(x_gm, [0, 0], [32, 128])
gamma = asctile.copy_in(gamma_gm, [0], [128])
beta = asctile.copy_in(beta_gm, [0], [128])
output, mean, rstd = asctile.layer_norm(input, gamma, beta, 1e-5)

Apply LayerNorm to a 1D tensor with variance output:

input = asctile.copy_in(x_gm, [0], [128])
gamma = asctile.copy_in(gamma_gm, [0], [128])
beta = asctile.copy_in(beta_gm, [0], [128])
output, mean, variance = asctile.layer_norm(input, gamma, beta, 1e-6, output_rstd=False)