Repository: Project-MONAI/MONAI
File: monai/networks/blocks/selfattention.py (Python)

Code before:
            k = self.input_rearrange(self.to_k(x))
            v = self.input_rearrange(self.to_v(x))

        if self.attention_dtype is not None:
            q = q.to(self.attention_dtype)
            k = k.to(self.attention_dtype)

        if self.use_flash_attention:
            x = F.scaled_dot_product_attention(
                query=q,
                key=k,
                value=v,
                attn_mask=attn_mask,
                scale=self.scale,
                dropout_p=self.dropout_rate,
                is_causal=self.causal,
            )
        else:
            att_mat = torch.einsum("blxd,blyd->blxy", q, k) * self.scale

            # apply relative positional embedding if defined
            if self.rel_positional_embedding is not None:
                att_mat = self.rel_positional_embedding(x, att_mat, q)

            if self.causal:
                assert attn_mask is None, "Causal attention does not support attention masks."
                att_mat = att_mat.masked_fill(self.causal_mask[:, :, : x.shape[-2], : x.shape[-2]] == 0, float("-inf"))

            if attn_mask is not None:
                attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
                attn_mask = attn_mask.expand(-1, self.num_heads, -1, -1)
                att_mat = att_mat.masked_fill(attn_mask == 0, float("-inf"))

            att_mat = att_mat.softmax(dim=-1)
            if self.save_attn:
                # no gradients and new tensor;
                # https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html
                self.att_mat = att_mat.detach()

            att_mat = self.drop_weights(att_mat)
            x = torch.einsum("bhxy,bhyd->bhxd", att_mat, v)

        x = self.out_rearrange(x)
        if self.include_fc:
            x = self.out_proj(x)
        x = self.drop_output(x)
        return x


Review comment (suggestion):
```suggestion
                if attn_mask is not None:
                    raise ValueError("Causal attention does not support attention masks.")
```

Code after:
            k = self.input_rearrange(self.to_k(x))
            v = self.input_rearrange(self.to_v(x))

        if self.attention_dtype is not None:
            q = q.to(self.attention_dtype)
            k = k.to(self.attention_dtype)

        if self.use_flash_attention:
            x = F.scaled_dot_product_attention(
                query=q,
                key=k,
                value=v,
                attn_mask=attn_mask,
                scale=self.scale,
                dropout_p=self.dropout_rate,
                is_causal=self.causal,
            )
        else:
            att_mat = torch.einsum("blxd,blyd->blxy", q, k) * self.scale

            # apply relative positional embedding if defined
            if self.rel_positional_embedding is not None:
                att_mat = self.rel_positional_embedding(x, att_mat, q)

            if self.causal:
                if attn_mask is not None:
                    raise ValueError("Causal attention does not support attention masks.")
                att_mat = att_mat.masked_fill(self.causal_mask[:, :, : x.shape[-2], : x.shape[-2]] == 0, float("-inf"))

            if attn_mask is not None:
                attn_mask = attn_mask.unsqueeze(1).unsqueeze(2)
                attn_mask = attn_mask.expand(-1, self.num_heads, -1, -1)
                att_mat = att_mat.masked_fill(attn_mask == 0, float("-inf"))

            att_mat = att_mat.softmax(dim=-1)
            if self.save_attn:
                # no gradients and new tensor;
                # https://pytorch.org/docs/stable/generated/torch.Tensor.detach.html
                self.att_mat = att_mat.detach()

            att_mat = self.drop_weights(att_mat)
            x = torch.einsum("bhxy,bhyd->bhxd", att_mat, v)

        x = self.out_rearrange(x)
        if self.include_fc:
            x = self.out_proj(x)
        x = self.drop_output(x)
        return x