Repository: huggingface/transformers
Pull request: Update get device properties and types

By using the more consistent `tuple[Optional[str], Optional[int], Optional[int]]` format @remi-or introduced as the functional type of the class we can simplify a lot of code.
We keep the old `tuple[Optional[str], Union[None, int, tuple[int, int]]]` as the internal type of `Expectations` because it helps us create instances of the class elegantly.

Changed files:
filename: src/transformers/testing_utils.py
patch: @@ -3223,7 +3223,9 @@ def cleanup(device: str, gc_collect=False):
 
 
 # Type definition of key used in `Expectations` class.
-DeviceProperties = tuple[Union[str, None], Union[None, int, tuple[int, int]]]
+DeviceProperties = tuple[Optional[str], Optional[int], Optional[int]]
+# Helper type. Makes creating instances of `Expectations` smoother.
+PackedDeviceProperties = tuple[Optional[str], Union[None, int, tuple[int, int]]]
 
 
 @cache
@@ -3236,57 +3238,61 @@ def get_device_properties() -> DeviceProperties:
 
         major, minor = torch.cuda.get_device_capability()
         if IS_ROCM_SYSTEM:
-            return ("rocm", (major, minor))
+            return ("rocm", major, minor)
         else:
-            return ("cuda", (major, minor))
+            return ("cuda", major, minor)
     elif IS_XPU_SYSTEM:
         import torch
 
         # To get more info of the architecture meaning and bit allocation, refer to https://github.com/intel/llvm/blob/sycl/sycl/include/sycl/ext/oneapi/experimental/device_architecture.def
         arch = torch.xpu.get_device_capability()["architecture"]
         gen_mask = 0x000000FF00000000
         gen = (arch & gen_mask) >> 32
-        return ("xpu", gen)
+        return ("xpu", gen, None)
     else:
-        return (torch_device, None)
+        return (torch_device, None, None)
 
 
-class Expectations(UserDict[DeviceProperties, Any]):
+class Expectations(UserDict[PackedDeviceProperties, Any]):
     def get_expectation(self) -> Any:
         """
         Find best matching expectation based on environment device properties.
         """
         return self.find_expectation(get_device_properties())
 
+    def unpacked(self) -> list[tuple[DeviceProperties, Any]]:
+        return [(Expectations.unpack(k), v) for k, v in self.data.items()]
+
     @staticmethod
-    def unpack(key: DeviceProperties) -> tuple[Optional[str], Optional[int], Optional[int]]:
+    def unpack(properties: PackedDeviceProperties) -> DeviceProperties:
         """
-        Unpack a `DeviceProperties` tuple into a `device_type`, `major`, and `minor`.
+        Unpack `PackedDeviceProperties` tuple into consistently formatted `DeviceProperties` tuple.
         """
-        device_type, major_minor = key
+        device_type, major_minor = properties
         if major_minor is None:
             major, minor = None, None
         elif isinstance(major_minor, int):
             major, minor = major_minor, None
         else:
             major, minor = major_minor
+
         return device_type, major, minor
 
     @staticmethod
-    def is_default(key: DeviceProperties) -> bool:
-        return all(p is None for p in key)
+    def is_default(properties: DeviceProperties) -> bool:
+        return all(p is None for p in properties)
 
     @staticmethod
-    def score(key: DeviceProperties, other: DeviceProperties) -> int:
+    def score(properties: DeviceProperties, other: DeviceProperties) -> float:
         """
         Returns score indicating how similar two instances of the `Properties` tuple are.
         Rules are as follows:
             * Matching `type` adds one point, semi-matching `type` adds half a point (e.g. cuda and rocm).
             * If types match, matching `major` adds another point, and then matching `minor` adds another.
             * Default expectation (if present) is worth 0.1 point to distinguish it from a straight-up zero.
         """
-        device_type, major, minor = Expectations.unpack(key)
-        other_device_type, other_major, other_minor = Expectations.unpack(other)
+        device_type, major, minor = properties
+        other_device_type, other_major, other_minor = other
 
         score = 0
         # Matching device type, maybe major and minor
@@ -3299,19 +3305,21 @@ def score(key: DeviceProperties, other: DeviceProperties) -> int:
         # Semi-matching device type
         elif device_type in ["cuda", "rocm"] and other_device_type in ["cuda", "rocm"]:
             score = 0.5
+
         # Default expectation
         if Expectations.is_default(other):
             score = 0.1
+
         return score
 
-    def find_expectation(self, key: DeviceProperties = (None, None)) -> Any:
+    def find_expectation(self, properties: DeviceProperties = (None, None, None)) -> Any:
         """
         Find best matching expectation based on provided device properties.
         """
-        (result_key, result) = max(self.data.items(), key=lambda x: Expectations.score(key, x[0]))
+        (result_key, result) = max(self.unpacked(), key=lambda x: Expectations.score(properties, x[0]))
 
-        if Expectations.score(key, result_key) == 0:
-            raise ValueError(f"No matching expectation found for {key}")
+        if Expectations.score(properties, result_key) == 0:
+            raise ValueError(f"No matching expectation found for {properties}")
 
         return result
 

filename: tests/models/musicgen/test_modeling_musicgen.py
patch: @@ -31,7 +31,6 @@
     T5Config,
 )
 from transformers.testing_utils import (
-    Expectations,
     get_device_properties,
     is_torch_available,
     require_flash_attn,
@@ -1052,8 +1051,7 @@ def test_sdpa_can_dispatch_on_flash(self):
         if not self.has_attentions:
             self.skipTest(reason="Model architecture does not support attentions")
 
-        device_properties = get_device_properties()
-        device_type, major, minor = Expectations.unpack(device_properties)
+        device_type, major, _ = get_device_properties()
         if device_type == "cuda" and major < 8:
             self.skipTest(reason="This test requires an NVIDIA GPU with compute capability >= 8.0")
         elif device_type == "rocm" and major < 9:

filename: tests/models/musicgen_melody/test_modeling_musicgen_melody.py
patch: @@ -30,7 +30,6 @@
     T5Config,
 )
 from transformers.testing_utils import (
-    Expectations,
     get_device_properties,
     is_torch_available,
     is_torchaudio_available,
@@ -1052,8 +1051,7 @@ def test_sdpa_can_dispatch_on_flash(self):
         if not self.has_attentions:
             self.skipTest(reason="Model architecture does not support attentions")
 
-        device_properties = get_device_properties()
-        device_type, major, minor = Expectations.unpack(device_properties)
+        device_type, major, _ = get_device_properties()
         if device_type == "cuda" and major < 8:
             self.skipTest(reason="This test requires an NVIDIA GPU with compute capability >= 8.0")
         elif device_type == "rocm" and major < 9:

filename: tests/test_modeling_common.py
patch: @@ -73,7 +73,6 @@
 )
 from transformers.testing_utils import (
     CaptureLogger,
-    Expectations,
     backend_empty_cache,
     get_device_properties,
     hub_retry,
@@ -3836,8 +3835,7 @@ def test_sdpa_can_dispatch_on_flash(self):
         if not self.has_attentions:
             self.skipTest(reason="Model architecture does not support attentions")
 
-        device_properties = get_device_properties()
-        device_type, major, minor = Expectations.unpack(device_properties)
+        device_type, major, _ = get_device_properties()
         if device_type == "cuda" and major < 8:
             self.skipTest(reason="This test requires an NVIDIA GPU with compute capability >= 8.0")
         elif device_type == "rocm" and major < 9:
@@ -3885,8 +3883,7 @@ def test_sdpa_can_compile_dynamic(self):
         if not self.has_attentions:
             self.skipTest(reason="Model architecture does not support attentions")
 
-        device_properties = get_device_properties()
-        device_type, major, minor = Expectations.unpack(device_properties)
+        device_type, major, _ = get_device_properties()
         if device_type == "cuda" and major < 8:
             self.skipTest(reason="This test requires an NVIDIA GPU with compute capability >= 8.0")
         elif device_type == "rocm" and major < 9: