    def __init__(
        self,
        in_features: int,
        hidden_features: Optional[int] = None,
        out_features: Optional[int] = None,
        act_layer: Module = GELU,
        dropout_rate: float = 0.0,
        bias_on: bool = True,
    ) -> None:
        """
        Args:
            in_features (int): Input feature dimension.
            hidden_features (Optional[int]): Hidden feature dimension. By default,
                hidden feature is set to input feature dimension.
            out_features (Optional[int]): Output feature dimension. By default, output
                features dimension is set to input feature dimension.
            act_layer (Callable): Activation layer used after the first linear layer.
            dropout_rate (float): Dropout rate after each linear layer. Dropout is not used
                by default.
        """
        super().__init__()
        self.dropout_rate = dropout_rate
        out_features = out_features or in_features
        hidden_features = hidden_features or in_features

        # TODO fc1 bias is set to zeros; unset if bias_on is True

        self.fc1 = Linear(
            in_features,
            hidden_features,
            bias=bias_on,
        )
        self.act = act_layer()
        self.fc2 = Linear(hidden_features, out_features, bias=bias_on)

        if self.dropout_rate > 0.0:
            self.dropout = Dropout(dropout_rate)
        else:
            self.dropout = Identity()
Args:
in_features (int): Input feature dimension.
hidden_features (Optional[int]): Hidden feature dimension. By default,
hidden feature is set to input feature dimension.
out_features (Optional[int]): Output feature dimension. By default, output
features dimension is set to input feature dimension.
act_layer (Callable): Activation layer used after the first linear layer.
dropout_rate (float): Dropout rate after each linear layer. Dropout is not used
by default.
python/aitemplate/frontend/nn/multiscale_attention.py