def METHOD_NAME(
    self,
    z: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    block_size=None,
) -> torch.Tensor:
    mask = mask.unsqueeze(-1)
    if not self._alphafold_original_mode:
        # divided by 1/sqrt(dim) for numerical stability
        mask = mask * (mask.shape[-2]**-0.5)
    z = self.layer_norm_in(z)
    if not self.training and block_size is not None:
        return self._chunk_2d(z, mask, block_size=block_size)
    g = nn.functional.linear(z, self.linear_g.weight)
    if self.training:
        ab = self.linear_ab_p(z) * mask * torch.sigmoid(
            self.linear_ab_g(z))
    else:
        ab = self.linear_ab_p(z)
        ab *= mask
        ab *= torch.sigmoid(self.linear_ab_g(z))
    a, b = torch.chunk(ab, 2, dim=-1)
    del z, ab
    if self.outgoing:
        a = permute_final_dims(a, (2, 0, 1))
        b = b.transpose(-1, -3)
    else:
        b = permute_final_dims(b, (2, 0, 1))
        a = a.transpose(-1, -3)
    x = torch.matmul(a, b)
    del a, b
    x = permute_final_dims(x, (1, 2, 0))
    x = self.layer_norm_out(x)
    x = nn.functional.linear(x, self.linear_z.weight)
    return x, g