Repository: Project-MONAI/MONAI
File: tests/test_crossattention.py (Python)

Code before:
    for hidden_size in [360, 480, 600, 768]:
        for num_heads in [4, 6, 8, 12]:
            for rel_pos_embedding in [None, RelPosEmbedding.DECOMPOSED]:
                for input_size in [(16, 32), (8, 8, 8)]:
                    for flash_attn in [True, False]:
                        test_case = [
                            {
                                "hidden_size": hidden_size,
                                "num_heads": num_heads,
                                "dropout_rate": dropout_rate,
                                "rel_pos_embedding": rel_pos_embedding if not flash_attn else None,
                                "input_size": input_size,
                                "use_flash_attention": flash_attn,
                            },
                            (2, 512, hidden_size),
                            (2, 512, hidden_size),
                        ]
                        TEST_CASE_CABLOCK.append(test_case)


class TestResBlock(unittest.TestCase):

    @parameterized.expand(TEST_CASE_CABLOCK)
    @skipUnless(has_einops, "Requires einops")
    def test_shape(self, input_param, input_shape, expected_shape):
        # Without flash attention
        net = CrossAttentionBlock(**input_param)
        with eval_mode(net):
            result = net(torch.randn(input_shape), context=torch.randn(2, 512, input_param["hidden_size"]))
            self.assertEqual(result.shape, expected_shape)

    def test_ill_arg(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(hidden_size=128, num_heads=12, dropout_rate=6.0)

        with self.assertRaises(ValueError):
            CrossAttentionBlock(hidden_size=620, num_heads=8, dropout_rate=0.4)

    @SkipIfBeforePyTorchVersion((2, 0))
    def test_save_attn_with_flash_attention(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(
                hidden_size=128, num_heads=3, dropout_rate=0.1, use_flash_attention=True, save_attn=True
            )

    @SkipIfBeforePyTorchVersion((2, 0))
    def test_rel_pos_embedding_with_flash_attention(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(
                hidden_size=128,
                num_heads=3,

Review comment (suggestion):
```suggestion
    @SkipIfBeforePyTorchVersion((2, 0))
    def test_shape(self, input_param, input_shape, expected_shape):
```

Code after:
    for hidden_size in [360, 480, 600, 768]:
        for num_heads in [4, 6, 8, 12]:
            for rel_pos_embedding in [None, RelPosEmbedding.DECOMPOSED]:
                for input_size in [(16, 32), (8, 8, 8)]:
                    for flash_attn in [True, False]:
                        test_case = [
                            {
                                "hidden_size": hidden_size,
                                "num_heads": num_heads,
                                "dropout_rate": dropout_rate,
                                "rel_pos_embedding": rel_pos_embedding if not flash_attn else None,
                                "input_size": input_size,
                                "use_flash_attention": flash_attn,
                            },
                            (2, 512, hidden_size),
                            (2, 512, hidden_size),
                        ]
                        TEST_CASE_CABLOCK.append(test_case)


class TestResBlock(unittest.TestCase):

    @parameterized.expand(TEST_CASE_CABLOCK)
    @skipUnless(has_einops, "Requires einops")
    @SkipIfBeforePyTorchVersion((2, 0))
    def test_shape(self, input_param, input_shape, expected_shape):
        # Without flash attention
        net = CrossAttentionBlock(**input_param)
        with eval_mode(net):
            result = net(torch.randn(input_shape), context=torch.randn(2, 512, input_param["hidden_size"]))
            self.assertEqual(result.shape, expected_shape)

    def test_ill_arg(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(hidden_size=128, num_heads=12, dropout_rate=6.0)

        with self.assertRaises(ValueError):
            CrossAttentionBlock(hidden_size=620, num_heads=8, dropout_rate=0.4)

    @SkipIfBeforePyTorchVersion((2, 0))
    def test_save_attn_with_flash_attention(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(
                hidden_size=128, num_heads=3, dropout_rate=0.1, use_flash_attention=True, save_attn=True
            )

    @SkipIfBeforePyTorchVersion((2, 0))
    def test_rel_pos_embedding_with_flash_attention(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(
                hidden_size=128,