def METHOD_NAME(
    self,
    z: torch.Tensor,
    mask: Optional[torch.Tensor] = None,
    block_size: int = None,
) -> torch.Tensor:
    # avoid too small chunk size
    # block_size = max(block_size, 256)
    new_z = z.new_zeros(z.shape)
    dim1 = z.shape[-3]
    def _slice_linear(z, linear: Linear, a=True):
        d_hid = linear.bias.shape[0] // 2
        index = 0 if a else d_hid
        p = (
            nn.functional.linear(z, linear.weight[index:index + d_hid])
            + linear.bias[index:index + d_hid])
        return p
    def _chunk_projection(z, mask, a=True):
        p = _slice_linear(z, self.linear_ab_p, a) * mask
        p *= torch.sigmoid(_slice_linear(z, self.linear_ab_g, a))
        return p
    num_chunk = (dim1 + block_size - 1) // block_size
    for i in range(num_chunk):
        chunk_start = i * block_size
        chunk_end = min(chunk_start + block_size, dim1)
        if self.outgoing:
            a_chunk = _chunk_projection(
                z[..., chunk_start:chunk_end, :, :],
                mask[..., chunk_start:chunk_end, :, :],
                a=True,
            )
            a_chunk = permute_final_dims(a_chunk, (2, 0, 1))
        else:
            a_chunk = _chunk_projection(
                z[..., :, chunk_start:chunk_end, :],
                mask[..., :, chunk_start:chunk_end, :],
                a=True,
            )
            a_chunk = a_chunk.transpose(-1, -3)
        for j in range(num_chunk):
            j_chunk_start = j * block_size
            j_chunk_end = min(j_chunk_start + block_size, dim1)
            if self.outgoing:
                b_chunk = _chunk_projection(
                    z[..., j_chunk_start:j_chunk_end, :, :],
                    mask[..., j_chunk_start:j_chunk_end, :, :],
                    a=False,
                )
                b_chunk = b_chunk.transpose(-1, -3)
            else:
                b_chunk = _chunk_projection(
                    z[..., :, j_chunk_start:j_chunk_end, :],
                    mask[..., :, j_chunk_start:j_chunk_end, :],
                    a=False,
                )
                b_chunk = permute_final_dims(b_chunk, (2, 0, 1))
            x_chunk = torch.matmul(a_chunk, b_chunk)
            del b_chunk
            x_chunk = permute_final_dims(x_chunk, (1, 2, 0))
            x_chunk = self.layer_norm_out(x_chunk)
            x_chunk = self.linear_z(x_chunk)
            x_chunk *= torch.sigmoid(
                self.linear_g(z[..., chunk_start:chunk_end,
                                j_chunk_start:j_chunk_end, :]))
            new_z[..., chunk_start:chunk_end,
                  j_chunk_start:j_chunk_end, :] = x_chunk
            del x_chunk
        del a_chunk
    return new_z