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

Code before:
                    test_case = [
                        {
                            "hidden_size": hidden_size,
                            "num_heads": num_heads,
                            "dropout_rate": dropout_rate,
                            "rel_pos_embedding": rel_pos_embedding,
                            "input_size": input_size,
                        },
                        (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)
        # With flash attention
        net = CrossAttentionBlock(**input_param, use_flash_attention=True)
        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((1, 13))
    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
            )

    @skipUnless(has_einops, "Requires einops")
    def test_attention_dim_not_multiple_of_heads(self):
        with self.assertRaises(ValueError):
            CrossAttentionBlock(hidden_size=128, num_heads=3, dropout_rate=0.1)

    @skipUnless(has_einops, "Requires einops")
    def test_inner_dim_different(self):

Review comment (refactor):
The alternative to this duplicate code is to add a `use_flash_attention` value to the `input_param` dictionary constructed above in the nested for-loop.

Code after:
                    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,
                num_heads=3,
                dropout_rate=0.1,
                use_flash_attention=True,
                save_attn=False,