Repository: huggingface/transformers
File: tests/kernels/test_kernels.py (Python)

Code before:
            load_and_register_attn_kernel(attn_impl)
            self.assertIn(attn_impl, ALL_ATTENTION_FUNCTIONS.valid_keys())
            # Cleanup registration to avoid leaking functions across tests
            try:
                ALL_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_ATTENTION_FUNCTIONS`: {e}")
            try:
                ALL_MASK_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up ALL_MASK_ATTENTION_FUNCTIONS: {e}")

    def test_load_and_register_named_function_kernel(self):
        def my_attention(*args, **kwargs):
            return None

        kernel_obj = types.SimpleNamespace(my_func=my_attention)
        with patch("transformers.integrations.hub_kernels.get_kernel", return_value=kernel_obj):
            attn_impl = "org/model:my_func"
            load_and_register_attn_kernel(attn_impl)
            self.assertIn(attn_impl, ALL_ATTENTION_FUNCTIONS.valid_keys())
            # Cleanup registration to avoid leaking functions across tests
            try:
                ALL_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up ALL_ATTENTION_FUNCTIONS: {e}")
            try:
                ALL_MASK_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up ALL_MASK_ATTENTION_FUNCTIONS: {e}")


@require_kernels
class TestUseKernelsLifecycle(TestCasePlus):
    @classmethod
    def setUpClass(cls):
        cls.model_id = "unsloth/Llama-3.2-1B-Instruct"
        cls.model = AutoModelForCausalLM.from_pretrained(cls.model_id, use_kernels=False, device_map=torch_device)

    @classmethod
    def tearDownClass(cls):
        # Delete large objects to drop references early
        if hasattr(cls, "model"):
            try:
                del cls.model
            except Exception as e:
                print(f"Could not delete model: {e}")

    def tearDown(self):
        # Free accelerator memory/cache and trigger GC
        cleanup(torch_device, gc_collect=True)

Review comment (refactor):
```suggestion
                print(f"Could not clean up `ALL_ATTENTION_FUNCTIONS`: {e}")
```

Code after:
            load_and_register_attn_kernel(attn_impl)
            self.assertIn(attn_impl, ALL_ATTENTION_FUNCTIONS.valid_keys())
            # Cleanup registration to avoid leaking functions across tests
            try:
                ALL_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_ATTENTION_FUNCTIONS`: {e}")
            try:
                ALL_MASK_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_MASK_ATTENTION_FUNCTIONS`: {e}")

    def test_load_and_register_named_function_kernel(self):
        def my_attention(*args, **kwargs):
            return None

        kernel_obj = types.SimpleNamespace(my_func=my_attention)
        with patch("transformers.integrations.hub_kernels.get_kernel", return_value=kernel_obj):
            attn_impl = "org/model:my_func"
            load_and_register_attn_kernel(attn_impl)
            self.assertIn(attn_impl, ALL_ATTENTION_FUNCTIONS.valid_keys())
            # Cleanup registration to avoid leaking functions across tests
            try:
                ALL_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_ATTENTION_FUNCTIONS`: {e}")
            try:
                ALL_MASK_ATTENTION_FUNCTIONS.pop(attn_impl, None)
            except Exception as e:
                print(f"Could not clean up `ALL_MASK_ATTENTION_FUNCTIONS`: {e}")


@require_kernels
class TestUseKernelsLifecycle(TestCasePlus):
    @classmethod
    def setUpClass(cls):
        cls.model_id = "unsloth/Llama-3.2-1B-Instruct"
        cls.model = AutoModelForCausalLM.from_pretrained(cls.model_id, use_kernels=False, device_map=torch_device)

    @classmethod
    def tearDownClass(cls):
        # Delete large objects to drop references early
        if hasattr(cls, "model"):
            try:
                del cls.model
            except Exception as e:
                print(f"Could not delete model: {e}")

    def tearDown(self):
        # Free accelerator memory/cache and trigger GC
        cleanup(torch_device, gc_collect=True)