def METHOD_NAME(
    self,
    x: torch.Tensor,
    x_lengths: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """Calculate forward propagation.
    Args:
        x (Tensor): Input index tensor (B, T_text).
        x_lengths (Tensor): Length tensor (B,).
    Returns:
        Tensor: Encoded hidden representation (B, attention_dim, T_text).
        Tensor: Projected mean tensor (B, attention_dim, T_text).
        Tensor: Projected scale tensor (B, attention_dim, T_text).
        Tensor: Mask tensor for input tensor (B, 1, T_text).
    """
    x = self.emb(x) * math.sqrt(self.attention_dim)
    x_mask = (
        make_non_pad_mask(x_lengths)
        .to(
            device=x.device,
            dtype=x.dtype,
        )
        .unsqueeze(1)
    )
    # encoder assume the channel last (B, T_text, attention_dim)
    # but mask shape shoud be (B, 1, T_text)
    x, _ = self.encoder(x, x_mask)
    # convert the channel first (B, attention_dim, T_text)
    x = x.transpose(1, 2)
    stats = self.proj(x) * x_mask
    m, logs = stats.split(stats.size(1) // 2, dim=1)
    return x, m, logs, x_mask