espnet2.speechlm.module.transformer.Linear

source

class espnet2.speechlm.module.transformer.Linear(in_features: int, out_features: int, bias: bool = True, device=None, dtype=None)

Bases: Linear

Linear layer with automatic dtype adjustment for inputs.

This class extends the PyTorch nn.Linear layer to ensure that the weights and biases are cast to the appropriate dtype of the input tensor x during the forward pass. This is particularly useful when working with mixed-precision training.

• Parameters:
  • in_features (int) – Number of input features.
  • out_features (int) – Number of output features.
  • bias (bool , optional) – If set to False, the layer will not learn an additive bias. Default is True.
• Returns: The output tensor after applying the linear transformation.
• Return type: Tensor

####### Examples

>>> linear_layer = Linear(10, 5)
>>> input_tensor = torch.randn(2, 10, dtype=torch.float32)
>>> output_tensor = linear_layer(input_tensor)
>>> output_tensor.shape
torch.Size([2, 5])

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

forward(x: Tensor) → Tensor

Computes the forward pass for the linear layer.

This method applies a linear transformation to the input tensor x using the layer’s weights and bias (if provided). The weights are cast to the same data type as the input tensor, ensuring compatibility for operations.

• Parameters:x (Tensor) – The input tensor of shape (N,

  *

  , in_features) where N is the batch size and * represents any number of additional dimensions.

• Returns: The output tensor of shape (N,

  *

  , out_features) where : out_features is the number of features in the output. The output tensor will have the same data type as the input tensor.

• Return type: Tensor

####### Examples

>>> linear_layer = Linear(5, 3)
>>> input_tensor = torch.randn(10, 5)
>>> output_tensor = linear_layer(input_tensor)
>>> output_tensor.shape
torch.Size([10, 3])

NOTE

The bias is optional; if not provided, the layer will perform the linear transformation without an additive bias term.