espnet2.speechlm.module.transformer.LayerNorm

source

class espnet2.speechlm.module.transformer.LayerNorm(normalized_shape: int | List[int] | Size, eps: float = 1e-05, elementwise_affine: bool = True, bias: bool = True, device=None, dtype=None)

Bases: LayerNorm

LayerNorm is a custom implementation of PyTorch’s LayerNorm that ensures the

input tensor is cast to float before applying layer normalization. This class inherits from the nn.LayerNorm module.

normalized_shape

The shape of the input tensor for normalization.

• Type: tuple

eps

A value added to the denominator for numerical stability.

• Type: float

elementwise_affine

Whether to include learnable parameters for scaling and shifting.

• Type: bool

• Parameters:

  • normalized_shape (int or tuple) – Input shape for normalization.
  • eps (float , optional) – A value added to the denominator for numerical stability (default: 1e-5).
  • elementwise_affine (bool , optional) – If True, learnable parameters are added (default: True).

• Returns: The layer-normalized output tensor.

• Return type: Tensor

####### Examples

>>> layer_norm = LayerNorm(normalized_shape=10)
>>> input_tensor = torch.randn(5, 10)
>>> output_tensor = layer_norm(input_tensor)
>>> output_tensor.shape
torch.Size([5, 10])

Initialize internal Module state, shared by both nn.Module and ScriptModule.

forward(x: Tensor) → Tensor

Applies Layer Normalization to the input tensor.

This method overrides the forward method of the nn.LayerNorm class to ensure that the input tensor is converted to float before applying layer normalization, and then the output is cast back to the original dtype of the input tensor.

• Parameters:x (Tensor) – The input tensor to be normalized. The tensor should be of type float, but the method will convert it to float for processing.
• Returns: The layer-normalized tensor, returned in the original dtype of : the input tensor.
• Return type: Tensor

####### Examples

>>> layer_norm = LayerNorm(normalized_shape=10)
>>> input_tensor = torch.randn(2, 10)
>>> output_tensor = layer_norm(input_tensor)
>>> output_tensor.dtype  # Should match the dtype of input_tensor
torch.float32