espnet2.asr.state_spaces.components.TransposedLN

source

class espnet2.asr.state_spaces.components.TransposedLN(d, scalar=True)

Bases: Module

Transposed LayerNorm module.

This module applies Layer Normalization over the second dimension of the input tensor, which is expected to have the shape (B, D, L), where B is the batch size, D is the feature dimension, and L can represent one or more length dimensions.

The implementation includes two modes:

1. Scalar mode: Normalizes the input using learned scale and shift parameters.
2. LayerNorm mode: Applies the standard LayerNorm to the input.

This module is currently not optimized for speed and a dedicated CUDA/Triton implementation could provide substantial performance improvements.

scalar

Indicates whether to use scalar parameters for normalization.

• Type: bool

• Parameters:

  • d (int) – The number of features in the input tensor.
  • scalar (bool , optional) – If True, use scalar parameters for normalization. Default is True.

• Returns: The normalized output tensor of the same shape as the input.

• Return type: Tensor

####### Examples

>>> layer_norm = TransposedLN(d=64)
>>> input_tensor = torch.randn(10, 64, 5)  # Batch of 10, 64 features, 5 length
>>> output_tensor = layer_norm(input_tensor)
>>> output_tensor.shape
torch.Size([10, 64, 5])  # Output retains the same shape

NOTE

This implementation is currently slow. Consider using a more optimized implementation for performance-critical applications.

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

forward(x)

Transposed LayerNorm module.

This module applies Layer Normalization over the second dimension (channels) of the input tensor. It assumes the input shape is (B, D, L), where B is the batch size, D is the number of features (channels), and L can be one or more additional dimensions.

This implementation may be slow for large inputs, and a dedicated CUDA/Triton implementation should provide substantial end-to-end speedup.

scalar

If True, uses scalar parameters for normalization. If False, uses standard LayerNorm.

• Type: bool

m

The mean parameter, used if scalar is True.

• Type: Parameter

s

The standard deviation parameter, used if scalar is True.

• Type: Parameter

ln

The LayerNorm module, used if scalar is False.

• Type:LayerNorm

• Parameters:

  • d (int) – The number of features (channels) to normalize.
  • scalar (bool) – If True, use scalar normalization; if False, use LayerNorm. Default: True.

• Returns: The normalized output tensor.

• Return type: Tensor

####### Examples

>>> layer_norm = TransposedLN(d=64)
>>> input_tensor = torch.randn(32, 64, 10)  # (B, D, L)
>>> output_tensor = layer_norm(input_tensor)

NOTE

This module is designed for use with transposed convolutional architectures where the channel dimension is the second dimension of the input tensor.