Repository: huggingface/transformers
Pull request: Follow-up typing checking fixes

# What does this PR do?

- Do proper type check in case jax is installed.
- Make sure older torch versions don't raise typing issues 


Changed files:
filename: src/transformers/generation/continuous_batching/requests.py
patch: @@ -46,13 +46,13 @@ def get_device_and_memory_breakdown() -> tuple[torch.device, int, int, int]:
         torch.xpu.empty_cache()
         torch.xpu.synchronize()
         total_memory = torch.xpu.get_device_properties(device).total_memory
-        reserved_memory = torch.xpu.memory_reserved(device)
-        allocated_memory = torch.xpu.memory_allocated(device)
+        reserved_memory = getattr(torch.xpu, "memory_reserved")(device)
+        allocated_memory = getattr(torch.xpu, "memory_allocated")(device)
     elif torch.backends.mps.is_available() and torch.backends.mps.is_built():
         device = torch.device("mps")
         # MPS memory reporting (PyTorch 2.0+)
         total_memory = torch.mps.driver_allocated_memory()
-        allocated_memory = total_memory - torch.mps.recommended_max_memory()
+        allocated_memory = total_memory - getattr(torch.mps, "recommended_max_memory")()
         reserved_memory = 0  # MPS does not track reserved separately
     else:
         device = torch.device("cpu")

filename: src/transformers/generation/utils.py
patch: @@ -3799,7 +3799,7 @@ def _prefill(
         else:
             # Even if we are not compiling the forward, flex is always compiled when used. With chunked prefill, we may
             # end up needing just a bit more graphs than the default (which is 8). Doing this avoids very cryptic warnings
-            torch._dynamo.config.cache_size_limit = 64
+            getattr(torch, "_dynamo").config.cache_size_limit = 64
 
             chunk_size = generation_config.prefill_chunk_size
             input_chunks = torch.split(input_ids, chunk_size, dim=-1)

filename: src/transformers/utils/import_utils.py
patch: @@ -212,7 +212,7 @@ def is_cuda_platform() -> bool:
     if is_torch_available():
         import torch
 
-        return torch.version.cuda is not None
+        return getattr(torch, "version").cuda is not None
     return False
 
 
@@ -221,7 +221,7 @@ def is_rocm_platform() -> bool:
     if is_torch_available():
         import torch
 
-        return torch.version.hip is not None
+        return getattr(torch, "version").hip is not None
     return False
 
 
@@ -503,9 +503,9 @@ def is_torch_bf16_gpu_available() -> bool:
         # Note: Emulated in software by Metal using fp32 for hardware without native support (like M1/M2)
         return torch.backends.mps.is_macos_or_newer(14, 0)
     if is_torch_musa_available():
-        return torch.musa.is_bf16_supported() if hasattr(torch, "musa") else False
+        return getattr(torch, "musa").is_bf16_supported()
     if is_torch_mlu_available():
-        return torch.mlu.is_bf16_supported() if hasattr(torch, "mlu") else False
+        return getattr(torch, "mlu").is_bf16_supported()
     if is_torch_neuron_available():
         return getattr(torch, "neuron").is_bf16_supported()
     return False
@@ -572,7 +572,8 @@ def is_torch_tf32_available() -> bool:
             if f"{device_info.major}{device_info.minor}" >= "22":
                 return True
         return False
-    if not torch.cuda.is_available() or torch.version.cuda is None:
+    torch_version = getattr(torch, "version")
+    if not torch.cuda.is_available() or torch_version.cuda is None:
         return False
     if torch.cuda.get_device_properties(torch.cuda.current_device()).major < 8:
         return False
@@ -920,9 +921,10 @@ def is_flash_attn_2_available() -> bool:
     import torch
 
     try:
-        if torch.version.cuda:
+        torch_version = getattr(torch, "version")
+        if torch_version.cuda:
             return version.parse(flash_attn_version) >= version.parse("2.1.0")
-        elif torch.version.hip:
+        elif torch_version.hip:
             # TODO: Bump the requirement to 2.1.0 once released in https://github.com/ROCmSoftwarePlatform/flash-attention
             return version.parse(flash_attn_version) >= version.parse("2.0.4")
         elif is_torch_mlu_available():
@@ -1339,7 +1341,7 @@ def is_torchdynamo_compiling() -> bool:
     try:
         import torch
 
-        return torch.compiler.is_compiling()
+        return getattr(torch, "compiler").is_compiling()
     except Exception:
         return False
 
@@ -1348,7 +1350,7 @@ def is_torchdynamo_exporting() -> bool:
     try:
         import torch
 
-        return torch.compiler.is_exporting()
+        return getattr(torch, "compiler").is_exporting()
     except Exception:
         return False
 
@@ -1366,7 +1368,7 @@ def is_fake_tensor(x) -> bool:
     try:
         import torch
 
-        return isinstance(x, torch._subclasses.FakeTensor)
+        return isinstance(x, getattr(torch, "_subclasses").FakeTensor)
     except Exception:
         return False
 
@@ -1394,7 +1396,7 @@ def is_jax_jitting(x):
     try:
         import jax
 
-        return isinstance(x.jax(), jax.core.Tracer)
+        return isinstance(x.jax(), getattr(jax, "core").Tracer)
     except Exception:
         return False
 

filename: src/transformers/utils/kernel_config.py
patch: @@ -48,7 +48,7 @@ def infer_device(model):
     dev_type = param.device.type
     if dev_type == "cuda":
         # Refine based on actual platform
-        if torch.version.hip is not None:
+        if getattr(torch, "version").hip is not None:
             return "rocm"
 
     return dev_type