From b74af4caf0510cba6979ad3b6ad3a171c279a9b8 Mon Sep 17 00:00:00 2001 From: Alvaro Moran Date: Thu, 8 Oct 2026 14:40:09 +0000 Subject: [PATCH] fix: detect TPU with torch.accelerator.current_accelerator() torch_tpu removed torch.backends.tpu (it only held allow_excess_precision, now torch.tpu.allow_excess_precision), so the TPU backend was detected as cpu and no torch-tpu variant could be loaded. current_accelerator() reports the "tpu" PrivateUse1 backend that torch_tpu registers on import, without requiring TPU hardware or initializing the runtime. Applied to the kernels backend detection and the noarch _ops.py template. --- .../src/pyproject/templates/torch/noarch/_ops.py | 7 ++----- kernels/src/kernels/backends.py | 8 ++------ 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py index e8d30b5ff..5e487141d 100644 --- a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py +++ b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py @@ -4,11 +4,8 @@ def get_backend() -> str: """Detect the backend by inspecting torch.""" import torch - if hasattr(torch.backends, "tpu"): - # torch_tpu sets torch.backends.tpu when it is imported (via - # torch's device-backend autoload), regardless of whether TPU - # hardware is present — analogous to torch.version.cuda being - # set on CUDA builds without a GPU. + acc = torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None + if acc is not None and acc.type == "tpu": return "tpu" elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron diff --git a/kernels/src/kernels/backends.py b/kernels/src/kernels/backends.py index cd4473421..b8e20a8a6 100644 --- a/kernels/src/kernels/backends.py +++ b/kernels/src/kernels/backends.py @@ -217,12 +217,8 @@ def _backend() -> Backend: if has_torch: import torch - if hasattr(torch.backends, "tpu"): - # torch_tpu sets torch.backends.tpu when it is imported (via - # torch's device-backend autoload), regardless of whether TPU - # hardware is present — analogous to torch.version.cuda being - # set on CUDA builds without a GPU. The hardware-gated - # torch.tpu namespace only appears on hosts with TPU devices. + acc = torch.accelerator.current_accelerator() if hasattr(torch, "accelerator") else None + if acc is not None and acc.type == "tpu": return TPU() elif hasattr(torch, "neuron"): # Needs to be sorted before specific Torch builds, since Neuron