def METHOD_NAME(self, xs_pad, mask):
    """Forward method.
    Args:
        xs_pad: (batch, time, size = n_heads * attn_dim)
        mask: (batch, 1, time), nonpadding is 1, padding is 0
    Returns:
        torch.Tensor: (batch, time, size)
    """
    batch_size, seq_len, _ = xs_pad.shape
    mixed_query_layer = self.query(xs_pad)  # (batch, time, size)
    mixed_key_layer = self.key(xs_pad)  # (batch, time, size)
    if mask is not None:
        mask = mask.eq(0)  # padding is 1, nonpadding is 0
    # (batch, n_heads, time)
    query_for_score = (
        self.query_att(mixed_query_layer).transpose(1, 2)
        / self.attention_head_size**0.5
    )
    if mask is not None:
        min_value = float(
            numpy.finfo(
                torch.tensor(0, dtype=query_for_score.dtype).numpy().dtype
            ).min
        )
        query_for_score = query_for_score.masked_fill(mask, min_value)
        query_weight = torch.softmax(query_for_score, dim=-1).masked_fill(mask, 0.0)
    else:
        query_weight = torch.softmax(query_for_score, dim=-1)
    query_weight = query_weight.unsqueeze(2)  # (batch, n_heads, 1, time)
    query_layer = self.transpose_for_scores(
        mixed_query_layer
    )  # (batch, n_heads, time, attn_dim)
    pooled_query = (
        torch.matmul(query_weight, query_layer)
        .transpose(1, 2)
        .reshape(-1, 1, self.num_attention_heads * self.attention_head_size)
    )  # (batch, 1, size = n_heads * attn_dim)
    pooled_query = self.dropout(pooled_query)
    pooled_query_repeat = pooled_query.repeat(1, seq_len, 1)  # (batch, time, size)
    mixed_query_key_layer = (
        mixed_key_layer * pooled_query_repeat
    )  # (batch, time, size)
    # (batch, n_heads, time)
    query_key_score = (
        self.key_att(mixed_query_key_layer) / self.attention_head_size**0.5
    ).transpose(1, 2)
    if mask is not None:
        min_value = float(
            numpy.finfo(
                torch.tensor(0, dtype=query_key_score.dtype).numpy().dtype
            ).min
        )
        query_key_score = query_key_score.masked_fill(mask, min_value)
        query_key_weight = torch.softmax(query_key_score, dim=-1).masked_fill(
            mask, 0.0
        )
    else:
        query_key_weight = torch.softmax(query_key_score, dim=-1)
    query_key_weight = query_key_weight.unsqueeze(2)  # (batch, n_heads, 1, time)
    key_layer = self.transpose_for_scores(
        mixed_query_key_layer
    )  # (batch, n_heads, time, attn_dim)
    pooled_key = torch.matmul(
        query_key_weight, key_layer
    )  # (batch, n_heads, 1, attn_dim)
    pooled_key = self.dropout(pooled_key)
    # NOTE: value = query, due to param sharing
    weighted_value = (pooled_key * query_layer).transpose(
        1, 2
    )  # (batch, time, n_heads, attn_dim)
    weighted_value = weighted_value.reshape(
        weighted_value.shape[:-2]
        + (self.num_attention_heads * self.attention_head_size,)
    )  # (batch, time, size)
    weighted_value = (
        self.dropout(self.transform(weighted_value)) + mixed_query_layer
    )
    return weighted_value