Module/Function Name: FusedDropoutLayerNorm
Class torch.nn.FusedDropoutLayerNorm(dim, dropout=0.1, eps=1e-5, elementwise_affine=True): """ Creates a fused dropout and layer normalization module. The dropout and layer normalization operations are performed together in a single layer.
Parameters:
- dim (int): Input dimension.
- dropout (float, optional): Dropout probability. Default: 0.1 (10% dropout).
- eps (float, optional): Epsilon value for layer normalization (std variance addition). Default: 1e-5.
- elementwise_affine (bool, optional): If True, provides learnable scaling and normalization weights. Default: True.
"""
def forward(x):
"""
Forward pass of the FusedDropoutLayerNorm module.
Parameters:
- x (Tensor): Input tensor to be processed.
Returns:
Tensor: Normalized and dropout-applied output tensor.
"""
x = self.dropout(x)
return self.layer_norm(x)
Example Usage:
Dim: 512
import torch
from torch import nn
x = torch.randn(1, 512)
model = nn.FusedDropoutLayerNorm(512)
out = model(x)
print(out.shape) # Output: torch.Size([1, 512])