From 9bd19371f12249e96ec929d231732e14e8075a9d Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 6 Jul 2026 14:26:05 +0000
Subject: [PATCH 01/13] kernels: add support for TPU

Make the Python loader recognize torch_tpu as a backend so a kernel
built on torch_tpu._internal.pallas.jax_op can be published to / fetched
from the Hub and dropped into kernelize(...) like any CUDA kernel.

- Add TPU backend dataclass and parse_backend.
- Probe torch_tpu via _get_torch_privateuse_backend_name() == "tpu",
  placed before the neuron/cuda/hip checks.
- Add a tpu section (torch_tpu, jax, libtpu) to both python_depends
  JSON files so per-backend dependency validation resolves.
- Add _TPURepos and 'tpu' to _validate_device_type's supported_devices
  so kernelize(model, device="tpu") works and resolves layer mappings.
- Add has_tpu / tpu_only markers in conftest and a test_tpu test that
  loads kernels-test/relu-tpu on a torch_tpu interpreter.
---
 kernels-data/src/python_dependencies.json | 14 ++++++++
 kernels/src/kernels/backends.py           | 25 +++++++++++++-
 kernels/src/kernels/layer/kernelize.py    |  4 +--
 kernels/src/kernels/layer/repos.py        | 22 ++++++++++++
 kernels/src/kernels/python_depends.json   | 14 ++++++++
 kernels/tests/conftest.py                 |  6 ++++
 kernels/tests/test_basic.py               |  7 ++++
 kernels/tests/test_layer.py               | 41 +++++++++++++++++++++++
 kernels/tests/test_variants.py            |  1 +
 9 files changed, 131 insertions(+), 3 deletions(-)

diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json
index 266904e8..a37bf8d6 100644
--- a/kernels-data/src/python_dependencies.json
+++ b/kernels-data/src/python_dependencies.json
@@ -34,6 +34,20 @@
         "nix": [],
         "python": [{ "pkg": "onednn-devel" }]
       }
+    },
+    "tpu": {
+      "torch_tpu": {
+        "nix": [],
+        "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
+      },
+      "jax": {
+        "nix": [],
+        "python": [{ "pkg": "jax", "import": "jax" }]
+      },
+      "libtpu": {
+        "nix": ["libtpu"],
+        "python": [{ "pkg": "libtpu" }]
+      }
     }
   }
 }
diff --git a/kernels/src/kernels/backends.py b/kernels/src/kernels/backends.py
index 5dff532a..3aa6fbec 100644
--- a/kernels/src/kernels/backends.py
+++ b/kernels/src/kernels/backends.py
@@ -127,6 +127,24 @@ def parse(s: str) -> "Neuron":
         return Neuron()
 
 
+@strict
+@dataclass(unsafe_hash=True)
+class TPU:
+    @property
+    def name(self) -> str:
+        return "tpu"
+
+    @property
+    def variant_str(self) -> str:
+        return "tpu"
+
+    @staticmethod
+    def parse(s: str) -> "TPU":
+        if s != "tpu":
+            raise ValueError(f"Invalid TPU variant string: {s!r}")
+        return TPU()
+
+
 @dataclass(unsafe_hash=True)
 class ROCm:
     _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"rocm(\d+)(\d+)")
@@ -179,6 +197,8 @@ def parse_backend(s: str) -> Backend:
         return Metal.parse(s)
     elif s == "neuron":
         return Neuron.parse(s)
+    elif s == "tpu":
+        return TPU.parse(s)
     elif s.startswith("cu"):
         return CUDA.parse(s)
     elif s.startswith("rocm"):
@@ -195,7 +215,10 @@ def _backend() -> Backend:
     if has_torch:
         import torch
 
-        if hasattr(torch, "neuron"):
+        if hasattr(torch, "tpu"):
+            # torch_tpu registers a torch.tpu namespace on import.
+            return TPU()
+        elif hasattr(torch, "neuron"):
             # Needs to be sorted before specific Torch builds, since Neuron
             # extension can be loaded into e.g. CUDA Torch builds.
             return Neuron()
diff --git a/kernels/src/kernels/layer/kernelize.py b/kernels/src/kernels/layer/kernelize.py
index 5294a5f5..60d221d6 100644
--- a/kernels/src/kernels/layer/kernelize.py
+++ b/kernels/src/kernels/layer/kernelize.py
@@ -192,7 +192,7 @@ def kernelize(
             `Mode.TRAINING | Mode.TORCH_COMPILE` kernelizes the model for training with
             `torch.compile`.
         device (`Union[str, torch.device]`, *optional*):
-            The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "xpu".
+            The device type to load kernels for. Supported device types are: "cuda", "mps", "npu", "rocm", "tpu", "xpu".
             The device type will be inferred from the model parameters when not provided.
         use_fallback (`bool`, *optional*, defaults to `True`):
             Whether to use the original forward method of modules when no compatible kernel could be found.
@@ -276,7 +276,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 def _validate_device_type(device_type: str) -> None:
     """Validate that the device type is supported."""
-    supported_devices = {"cpu", "cuda", "mps", "neuron", "npu", "rocm", "xpu"}
+    supported_devices = {"cpu", "cuda", "mps", "neuron", "npu", "rocm", "tpu", "xpu"}
     if device_type not in supported_devices:
         raise ValueError(
             f"Unsupported device type '{device_type}'. Supported device types are: {', '.join(sorted(supported_devices))}"
diff --git a/kernels/src/kernels/layer/repos.py b/kernels/src/kernels/layer/repos.py
index fd02707c..c53c833b 100644
--- a/kernels/src/kernels/layer/repos.py
+++ b/kernels/src/kernels/layer/repos.py
@@ -37,6 +37,8 @@ def create_repo(device: Device) -> "DeviceRepos":
             return _NPURepos()
         elif device.type == "neuron":
             return _NeuronRepos()
+        elif device.type == "tpu":
+            return _TPURepos()
         else:
             raise ValueError(f"Unknown device type: {device.type}")
 
@@ -114,6 +116,26 @@ def insert(self, device: Device, repos: dict[Mode, RepositoryProtocol]):
         self._repos = repos
 
 
+class _TPURepos(DeviceRepos):
+    _repos: dict[Mode, RepositoryProtocol]
+
+    def __init__(self):
+        super().__init__()
+        self._repos = {}
+
+    @property
+    def repos(
+        self,
+    ) -> dict[Mode, RepositoryProtocol] | None:
+        return self._repos
+
+    def insert(self, device: Device, repos: dict[Mode, RepositoryProtocol]):
+        if device.type != "tpu":
+            raise ValueError(f"Device type must be 'tpu', got {device.type}")
+
+        self._repos = repos
+
+
 class _NPURepos(DeviceRepos):
     _repos: dict[Mode, RepositoryProtocol]
 
diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json
index 266904e8..a37bf8d6 100644
--- a/kernels/src/kernels/python_depends.json
+++ b/kernels/src/kernels/python_depends.json
@@ -34,6 +34,20 @@
         "nix": [],
         "python": [{ "pkg": "onednn-devel" }]
       }
+    },
+    "tpu": {
+      "torch_tpu": {
+        "nix": [],
+        "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
+      },
+      "jax": {
+        "nix": [],
+        "python": [{ "pkg": "jax", "import": "jax" }]
+      },
+      "libtpu": {
+        "nix": ["libtpu"],
+        "python": [{ "pkg": "libtpu" }]
+      }
     }
   }
 }
diff --git a/kernels/tests/conftest.py b/kernels/tests/conftest.py
index a9cb3c01..513261ff 100644
--- a/kernels/tests/conftest.py
+++ b/kernels/tests/conftest.py
@@ -34,6 +34,8 @@
 
 has_npu = torch is not None and _get_torch_privateuse_backend_name() == "npu"
 
+has_tpu = torch is not None and hasattr(torch, "tpu") and torch.tpu.device_count() > 0
+
 has_jax = importlib.util.find_spec("jax") is not None and importlib.util.find_spec("jax_tvm_ffi") is not None
 
 
@@ -53,6 +55,8 @@ def device():
         return "xpu"
     elif has_npu:
         return "npu"
+    elif has_tpu:
+        return "tpu"
 
     return "cpu"
 
@@ -74,5 +78,7 @@ def pytest_runtest_setup(item):
         pytest.skip("skipping XPU-only test on host without XPU")
     if "npu_only" in item.keywords and not has_npu:
         pytest.skip("skipping NPU-only test on host without NPU")
+    if "tpu_only" in item.keywords and not has_tpu:
+        pytest.skip("skipping TPU-only test on host without TPU")
     if "token" in item.keywords and not item.config.getoption("--token"):
         pytest.skip("need --token option to run this test")
diff --git a/kernels/tests/test_basic.py b/kernels/tests/test_basic.py
index 1460f945..5c9ce8db 100644
--- a/kernels/tests/test_basic.py
+++ b/kernels/tests/test_basic.py
@@ -229,6 +229,13 @@ def test_neuron():
     torch.testing.assert_close(relu.relu(x), x.relu())
 
 
+@pytest.mark.tpu_only
+def test_tpu():
+    relu = get_kernel("kernels-test/relu-tpu", version=1)
+    x = torch.randn((16, 16), device="tpu")
+    torch.testing.assert_close(relu.relu(x), x.relu())
+
+
 def test_trust_remote_code_blocks_untrusted_org():
     """Kernels from untrusted orgs should be rejected by default."""
     with pytest.raises(ValueError, match=r"not from a trusted publisher"):
diff --git a/kernels/tests/test_layer.py b/kernels/tests/test_layer.py
index e67abfc2..790e6445 100644
--- a/kernels/tests/test_layer.py
+++ b/kernels/tests/test_layer.py
@@ -255,6 +255,47 @@ def forward(self, x):
     assert smol.relu.n_calls == 1
 
 
+@pytest.mark.tpu_only
+def test_hub_forward_tpu():
+    torch.manual_seed(0)
+
+    mapping = {"ReLU": {"tpu": LayerRepository(repo_id="kernels-test/relu-tpu", version=1, layer_name="ReLU")}}
+
+    relu = ReLU()
+    X = torch.randn((16, 16), device="tpu")
+    Y = relu(X)
+
+    with use_kernel_mapping(mapping):
+        relu_with_kernel = kernelize(ReLUWithKernel(), device="tpu", mode=Mode.INFERENCE)
+    Y_kernel = relu_with_kernel(X)
+
+    torch.testing.assert_close(Y_kernel, Y)
+
+    assert relu.n_calls == 1
+    assert relu_with_kernel.n_calls == 0
+
+    # Check that the device type can be determined automatically.
+    class SMOL(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = nn.Linear(16, 16)
+            self.relu = ReLUWithKernel()
+
+        def forward(self, x):
+            return self.relu(self.linear(x))
+
+    smol = SMOL().to("tpu")
+
+    Y = smol(X)
+
+    with use_kernel_mapping(mapping):
+        smol = kernelize(smol, mode=Mode.INFERENCE)
+    Y_kernel = smol(X)
+
+    torch.testing.assert_close(Y, Y_kernel)
+    assert smol.relu.n_calls == 1
+
+
 @pytest.mark.rocm_only
 def test_hub_forward_rocm():
     torch.manual_seed(0)
diff --git a/kernels/tests/test_variants.py b/kernels/tests/test_variants.py
index 008e2c8a..1707c720 100644
--- a/kernels/tests/test_variants.py
+++ b/kernels/tests/test_variants.py
@@ -74,6 +74,7 @@
     "torch-metal",
     "torch-neuron",
     "torch-rocm",
+    "torch-tpu",
     "torch-xpu",
     "torch-npu",
     "torch-universal",

From a780b4b6437dea2e6e10c0f57a8b19abb73cb35b Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 6 Jul 2026 14:42:20 +0000
Subject: [PATCH 02/13] kernels-data: add Backend::Tpu enum variant

Add Tpu to the canonical Backend enum in kernels-data so metadata-tpu.json
and the Python-level kernels_data.Backend.TPU can be parsed and emitted.
TPU is a noarch backend, so the noarch-only TpuGeneral block follows the
shape of NeuronGeneral.

- Add Backend::Tpu to mod.rs (and to Backend::all() / as_str / Display /
  FromStr).
- Add tpu: Option<TpuGeneral> to the canonical General struct, and to the
  v3/v4/v5 per-edition schemas with matching From<*> / From<super::*>
  conversions both directions.
- v2 migration seeds tpu: None.
- Add PyBackend::Tpu (binding name "TPU") to the Python class, with
  both From directions and the __repr__ match arm; update the .pyi
  stub.
- Add Backend.TPU member and update the Backend.from_str docstring in
  the .pyi stub.
- Update the metadata.rs unit-test General literals to include tpu: None.
---
 docs/source/kernel-requirements.md            |  2 +-
 kernels-data/bindings/python/kernels_data.pyi |  3 +-
 kernels-data/bindings/python/src/lib.rs       |  7 ++++-
 .../python/tests/test_kernels_data.py         |  5 +++-
 kernels-data/src/config/mod.rs                | 12 +++++++-
 kernels-data/src/config/v2.rs                 |  1 +
 kernels-data/src/config/v3.rs                 | 19 ++++++++++++
 kernels-data/src/config/v4.rs                 | 19 ++++++++++++
 kernels-data/src/config/v5.rs                 | 29 +++++++++++++++++++
 kernels-data/src/metadata.rs                  |  2 ++
 10 files changed, 94 insertions(+), 5 deletions(-)

diff --git a/docs/source/kernel-requirements.md b/docs/source/kernel-requirements.md
index eb432481..2c4dd34f 100644
--- a/docs/source/kernel-requirements.md
+++ b/docs/source/kernel-requirements.md
@@ -149,7 +149,7 @@ The `backend` specifies a dictionary of the following form:
 ```
 
 The backend `type` must be one of `cann`, `cpu`, `cuda`, `metal`, `neuron`,
-`rocm`, or `xpu`. For CUDA and ROCm, the supported architectures must
+`rocm`, `tpu`, or `xpu`. For CUDA and ROCm, the supported architectures must
 be specified in the `archs` field.
 
 ### Python dependencies
diff --git a/kernels-data/bindings/python/kernels_data.pyi b/kernels-data/bindings/python/kernels_data.pyi
index 20e29619..43b7af6a 100644
--- a/kernels-data/bindings/python/kernels_data.pyi
+++ b/kernels-data/bindings/python/kernels_data.pyi
@@ -32,6 +32,7 @@ class Backend(Enum):
     Metal = "Metal"
     Neuron = "Neuron"
     ROCm = "ROCm"
+    TPU = "TPU"
     XPU = "XPU"
 
     @staticmethod
@@ -40,7 +41,7 @@ class Backend(Enum):
 
         Args:
             s: One of `"cann"`, `"cpu"`, `"cuda"`, `"metal"`,
-               `"neuron"`, `"rocm"`, `"xpu"`.
+               `"neuron"`, `"rocm"`, `"tpu"`, `"xpu"`.
 
         Raises:
             ValueError: If the backend name is unknown.
diff --git a/kernels-data/bindings/python/src/lib.rs b/kernels-data/bindings/python/src/lib.rs
index 9aa29f06..3cfcae33 100644
--- a/kernels-data/bindings/python/src/lib.rs
+++ b/kernels-data/bindings/python/src/lib.rs
@@ -88,6 +88,8 @@ enum PyBackend {
     Neuron,
     #[pyo3(name = "ROCm")]
     Rocm,
+    #[pyo3(name = "TPU")]
+    Tpu,
     #[pyo3(name = "XPU")]
     Xpu,
 }
@@ -101,6 +103,7 @@ impl From<Backend> for PyBackend {
             Backend::Metal => PyBackend::Metal,
             Backend::Neuron => PyBackend::Neuron,
             Backend::Rocm => PyBackend::Rocm,
+            Backend::Tpu => PyBackend::Tpu,
             Backend::Xpu => PyBackend::Xpu,
         }
     }
@@ -115,6 +118,7 @@ impl From<PyBackend> for Backend {
             PyBackend::Metal => Backend::Metal,
             PyBackend::Neuron => Backend::Neuron,
             PyBackend::Rocm => Backend::Rocm,
+            PyBackend::Tpu => Backend::Tpu,
             PyBackend::Xpu => Backend::Xpu,
         }
     }
@@ -123,7 +127,7 @@ impl From<PyBackend> for Backend {
 #[pymethods]
 impl PyBackend {
     /// Parse a backend name (`"cann"`, `"cpu"`, `"cuda"`, `"metal"`,
-    /// `"neuron"`, `"rocm"`, `"xpu"`).
+    /// `"neuron"`, `"rocm"`, `"tpu"`, `"xpu"`).
     #[staticmethod]
     #[pyo3(name = "from_str")]
     fn py_from_str(s: &str) -> PyResult<Self> {
@@ -144,6 +148,7 @@ impl PyBackend {
             PyBackend::Metal => "Metal",
             PyBackend::Neuron => "Neuron",
             PyBackend::Rocm => "ROCm",
+            PyBackend::Tpu => "TPU",
             PyBackend::Xpu => "XPU",
         };
         format!("Backend.{variant}")
diff --git a/kernels-data/bindings/python/tests/test_kernels_data.py b/kernels-data/bindings/python/tests/test_kernels_data.py
index 29bb52fe..87f9d111 100644
--- a/kernels-data/bindings/python/tests/test_kernels_data.py
+++ b/kernels-data/bindings/python/tests/test_kernels_data.py
@@ -68,7 +68,7 @@ def test_backend_hash():
 
 def test_backend_unknown():
     with pytest.raises(ValueError):
-        Backend.from_str("tpu")
+        Backend.from_str("dsp")
 
 
 def test_backend_all_variants_and_casing():
@@ -80,8 +80,11 @@ def test_backend_all_variants_and_casing():
     assert repr(Backend.ROCm) == "Backend.ROCm"
     assert repr(Backend.XPU) == "Backend.XPU"
     assert repr(Backend.CANN) == "Backend.CANN"
+    assert str(Backend.TPU) == "tpu"
+    assert repr(Backend.TPU) == "Backend.TPU"
     assert Backend.from_str("cann") == Backend.CANN
     assert Backend.from_str("ROCM") == Backend.ROCm
+    assert Backend.from_str("TPU") == Backend.TPU
     assert Backend.from_str("metal") == Backend.Metal
 
 
diff --git a/kernels-data/src/config/mod.rs b/kernels-data/src/config/mod.rs
index bc7855f7..3300b48c 100644
--- a/kernels-data/src/config/mod.rs
+++ b/kernels-data/src/config/mod.rs
@@ -113,6 +113,7 @@ pub struct General {
 
     pub cuda: Option<CudaGeneral>,
     pub neuron: Option<NeuronGeneral>,
+    pub tpu: Option<TpuGeneral>,
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -191,6 +192,10 @@ pub struct NeuronGeneral {
     pub python_depends: Option<Vec<String>>,
 }
 
+pub struct TpuGeneral {
+    pub python_depends: Option<Vec<String>>,
+}
+
 pub struct Hub {
     pub repo_id: Option<String>,
     pub branch: Option<String>,
@@ -379,11 +384,12 @@ pub enum Backend {
     Metal,
     Neuron,
     Rocm,
+    Tpu,
     Xpu,
 }
 
 impl Backend {
-    pub const fn all() -> [Backend; 7] {
+    pub const fn all() -> [Backend; 8] {
         [
             Backend::Cann,
             Backend::Cpu,
@@ -391,6 +397,7 @@ impl Backend {
             Backend::Metal,
             Backend::Neuron,
             Backend::Rocm,
+            Backend::Tpu,
             Backend::Xpu,
         ]
     }
@@ -403,6 +410,7 @@ impl Backend {
             Backend::Metal => "metal",
             Backend::Neuron => "neuron",
             Backend::Rocm => "rocm",
+            Backend::Tpu => "tpu",
             Backend::Xpu => "xpu",
         }
     }
@@ -417,6 +425,7 @@ impl Display for Backend {
             Backend::Metal => write!(f, "metal"),
             Backend::Neuron => write!(f, "neuron"),
             Backend::Rocm => write!(f, "rocm"),
+            Backend::Tpu => write!(f, "tpu"),
             Backend::Xpu => write!(f, "xpu"),
         }
     }
@@ -433,6 +442,7 @@ impl FromStr for Backend {
             "metal" => Ok(Backend::Metal),
             "neuron" => Ok(Backend::Neuron),
             "rocm" => Ok(Backend::Rocm),
+            "tpu" => Ok(Backend::Tpu),
             "xpu" => Ok(Backend::Xpu),
             _ => Err(format!("Unknown backend: {s}")),
         }
diff --git a/kernels-data/src/config/v2.rs b/kernels-data/src/config/v2.rs
index eed543ba..f79204a8 100644
--- a/kernels-data/src/config/v2.rs
+++ b/kernels-data/src/config/v2.rs
@@ -179,6 +179,7 @@ impl General {
             hub: general.hub.map(Into::into),
             neuron: None,
             python_depends: None,
+            tpu: None,
             xpu: None,
         }
     }
diff --git a/kernels-data/src/config/v3.rs b/kernels-data/src/config/v3.rs
index 590d16e6..766a0cf9 100644
--- a/kernels-data/src/config/v3.rs
+++ b/kernels-data/src/config/v3.rs
@@ -50,6 +50,8 @@ pub struct General {
 
     pub python_depends: Option<Vec<String>>,
 
+    pub tpu: Option<TpuGeneral>,
+
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -67,6 +69,12 @@ pub struct NeuronGeneral {
     pub python_depends: Option<Vec<String>>,
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
+pub struct TpuGeneral {
+    pub python_depends: Option<Vec<String>>,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub struct XpuGeneral {
@@ -155,6 +163,7 @@ pub enum Backend {
     Metal,
     Neuron,
     Rocm,
+    Tpu,
     Xpu,
 }
 
@@ -205,6 +214,7 @@ impl TryFrom<General> for super::General {
             hub: general.hub.map(Into::into),
             neuron: general.neuron.map(Into::into),
             python_depends: general.python_depends,
+            tpu: general.tpu.map(Into::into),
             xpu: general.xpu.map(Into::into),
         })
     }
@@ -228,6 +238,14 @@ impl From<NeuronGeneral> for super::NeuronGeneral {
     }
 }
 
+impl From<TpuGeneral> for super::TpuGeneral {
+    fn from(tpu: TpuGeneral) -> Self {
+        Self {
+            python_depends: tpu.python_depends,
+        }
+    }
+}
+
 impl From<XpuGeneral> for super::XpuGeneral {
     fn from(xpu: XpuGeneral) -> Self {
         Self {
@@ -279,6 +297,7 @@ impl From<Backend> for super::Backend {
             Backend::Metal => super::Backend::Metal,
             Backend::Neuron => super::Backend::Neuron,
             Backend::Rocm => super::Backend::Rocm,
+            Backend::Tpu => super::Backend::Tpu,
             Backend::Xpu => super::Backend::Xpu,
         }
     }
diff --git a/kernels-data/src/config/v4.rs b/kernels-data/src/config/v4.rs
index a88db69c..ecc6c012 100644
--- a/kernels-data/src/config/v4.rs
+++ b/kernels-data/src/config/v4.rs
@@ -49,6 +49,8 @@ pub struct General {
 
     pub python_depends: Option<Vec<String>>,
 
+    pub tpu: Option<TpuGeneral>,
+
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -66,6 +68,12 @@ pub struct NeuronGeneral {
     pub python_depends: Option<Vec<String>>,
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
+pub struct TpuGeneral {
+    pub python_depends: Option<Vec<String>>,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub struct XpuGeneral {
@@ -171,6 +179,7 @@ pub enum Backend {
     Metal,
     Neuron,
     Rocm,
+    Tpu,
     Xpu,
 }
 
@@ -203,6 +212,7 @@ impl From<General> for super::General {
             hub: general.hub.map(Into::into),
             neuron: general.neuron.map(Into::into),
             python_depends: general.python_depends,
+            tpu: general.tpu.map(Into::into),
             xpu: general.xpu.map(Into::into),
         }
     }
@@ -238,6 +248,14 @@ impl From<NeuronGeneral> for super::NeuronGeneral {
     }
 }
 
+impl From<TpuGeneral> for super::TpuGeneral {
+    fn from(tpu: TpuGeneral) -> Self {
+        Self {
+            python_depends: tpu.python_depends,
+        }
+    }
+}
+
 impl From<XpuGeneral> for super::XpuGeneral {
     fn from(xpu: XpuGeneral) -> Self {
         Self {
@@ -300,6 +318,7 @@ impl From<Backend> for super::Backend {
             Backend::Metal => super::Backend::Metal,
             Backend::Neuron => super::Backend::Neuron,
             Backend::Rocm => super::Backend::Rocm,
+            Backend::Tpu => super::Backend::Tpu,
             Backend::Xpu => super::Backend::Xpu,
         }
     }
diff --git a/kernels-data/src/config/v5.rs b/kernels-data/src/config/v5.rs
index 0154eeef..3e8a9996 100644
--- a/kernels-data/src/config/v5.rs
+++ b/kernels-data/src/config/v5.rs
@@ -62,6 +62,8 @@ pub struct General {
 
     pub python_depends: Option<Vec<String>>,
 
+    pub tpu: Option<TpuGeneral>,
+
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -79,6 +81,12 @@ pub struct NeuronGeneral {
     pub python_depends: Option<Vec<String>>,
 }
 
+#[derive(Debug, Deserialize, Serialize)]
+#[serde(deny_unknown_fields, rename_all = "kebab-case")]
+pub struct TpuGeneral {
+    pub python_depends: Option<Vec<String>>,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 #[serde(deny_unknown_fields, rename_all = "kebab-case")]
 pub struct XpuGeneral {
@@ -185,6 +193,7 @@ pub enum Backend {
     Metal,
     Neuron,
     Rocm,
+    Tpu,
     Xpu,
 }
 
@@ -217,6 +226,7 @@ impl From<General> for super::General {
             hub: general.hub.map(Into::into),
             neuron: general.neuron.map(Into::into),
             python_depends: general.python_depends,
+            tpu: general.tpu.map(Into::into),
             xpu: general.xpu.map(Into::into),
         }
     }
@@ -252,6 +262,14 @@ impl From<NeuronGeneral> for super::NeuronGeneral {
     }
 }
 
+impl From<TpuGeneral> for super::TpuGeneral {
+    fn from(tpu: TpuGeneral) -> Self {
+        Self {
+            python_depends: tpu.python_depends,
+        }
+    }
+}
+
 impl From<XpuGeneral> for super::XpuGeneral {
     fn from(xpu: XpuGeneral) -> Self {
         Self {
@@ -313,6 +331,7 @@ impl From<Backend> for super::Backend {
             Backend::Metal => super::Backend::Metal,
             Backend::Neuron => super::Backend::Neuron,
             Backend::Rocm => super::Backend::Rocm,
+            Backend::Tpu => super::Backend::Tpu,
             Backend::Xpu => super::Backend::Xpu,
         }
     }
@@ -420,6 +439,7 @@ impl From<super::General> for General {
             hub: general.hub.map(Into::into),
             neuron: general.neuron.map(Into::into),
             python_depends: general.python_depends,
+            tpu: general.tpu.map(Into::into),
             xpu: general.xpu.map(Into::into),
         }
     }
@@ -455,6 +475,14 @@ impl From<super::NeuronGeneral> for NeuronGeneral {
     }
 }
 
+impl From<super::TpuGeneral> for TpuGeneral {
+    fn from(tpu: super::TpuGeneral) -> Self {
+        Self {
+            python_depends: tpu.python_depends,
+        }
+    }
+}
+
 impl From<super::XpuGeneral> for XpuGeneral {
     fn from(xpu: super::XpuGeneral) -> Self {
         Self {
@@ -516,6 +544,7 @@ impl From<super::Backend> for Backend {
             super::Backend::Metal => Backend::Metal,
             super::Backend::Neuron => Backend::Neuron,
             super::Backend::Rocm => Backend::Rocm,
+            super::Backend::Tpu => Backend::Tpu,
             super::Backend::Xpu => Backend::Xpu,
         }
     }
diff --git a/kernels-data/src/metadata.rs b/kernels-data/src/metadata.rs
index d895f7fb..3da04dc4 100644
--- a/kernels-data/src/metadata.rs
+++ b/kernels-data/src/metadata.rs
@@ -159,6 +159,7 @@ mod tests {
                 python_depends: None,
                 cuda: None,
                 neuron: None,
+                tpu: None,
                 xpu: None,
             },
             kernels: HashMap::new(),
@@ -214,6 +215,7 @@ mod tests {
                 python_depends: None,
                 cuda: None,
                 neuron: None,
+                tpu: None,
                 xpu: None,
             },
             kernels: HashMap::new(),

From dac23cd5956e296c42319db7cc1a1f7a9fc1ab63 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 6 Jul 2026 14:46:55 +0000
Subject: [PATCH 03/13] kernel-builder: detect TPU in noarch _ops.py template

Add a 'tpu' arm to the templated get_backend() function, placed before
the existing torch.cuda/hip/xpu probes and matching the canonical
privateuse1 backend name check used in kernels/backends.py and
kernels/tests/conftest.py. The three layers (_backend(), has_tpu,
_kernel._ops.get_backend()) must agree so a noarch kernel built with
backends=["tpu"] finds the matching torch_tpu ops namespace at import
time.
---
 docs/source/builder-cli.md                                  | 2 +-
 kernel-builder/src/init.rs                                  | 2 +-
 kernel-builder/src/pyproject/templates/torch/noarch/_ops.py | 5 ++++-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/docs/source/builder-cli.md b/docs/source/builder-cli.md
index 3fd7d764..532bc8ba 100644
--- a/docs/source/builder-cli.md
+++ b/docs/source/builder-cli.md
@@ -83,7 +83,7 @@ Initialize a new kernel project from template
 
   Default value: `Apache-2.0`
 * `--name <OWNER/REPO>` — Name of the kernel repo (e.g. `drbh/my-kernel`)
-* `--backends <BACKENDS>` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`)
+* `--backends <BACKENDS>` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `tpu`, `xpu`)
 
 * `--overwrite` — Overwrite existing scaffold files (preserves other files)
 
diff --git a/kernel-builder/src/init.rs b/kernel-builder/src/init.rs
index 8ce37f48..f12b4edd 100644
--- a/kernel-builder/src/init.rs
+++ b/kernel-builder/src/init.rs
@@ -44,7 +44,7 @@ pub struct InitArgs {
     #[arg(long, value_name = "OWNER/REPO")]
     pub name: Option<RepoInfo>,
 
-    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`).
+    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `tpu`, `xpu`).
     #[arg(long, num_args = 1.., default_values_t = default_init_backends())]
     pub backends: Vec<BackendSelection>,
 
diff --git a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py
index 75fd8e08..46c697e9 100644
--- a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py
+++ b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py
@@ -4,7 +4,10 @@ def get_backend() -> str:
     """Detect the backend by inspecting torch."""
     import torch
 
-    if hasattr(torch, "neuron"):
+    if hasattr(torch, "tpu"):
+        # torch_tpu registers a torch.tpu namespace on import.
+        return "tpu"
+    elif hasattr(torch, "neuron"):
         # Needs to be sorted before specific Torch builds, since Neuron
         # extension can be loaded into e.g. CUDA Torch builds.
         return "neuron"

From 31c394ba0aced82a52033e1ac17dc59c89d95b5c Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 6 Jul 2026 14:53:46 +0000
Subject: [PATCH 04/13] examples: add kernels-test/relu-tpu example kernel
 repository

A noarch Pallas-backed ReLU kernel that serves as the reference for
TPU kernel authors. Built on torch_tpu._internal.pallas.jax_op, so the
op namespace discovered at import is just the torch_tpu one registered
by the Torch XLA plugin rather than the kernel-builder _OPS_NAME
namespace (which is still emitted for symmetry but unused here).

- build.toml edition 5 with backends=["tpu"], [general.tpu]
  python-depends = ["torch_tpu", "jax"], [general.hub] repo-id =
  "kernels-test/relu-tpu", and no per-backend kernel sources
  (torch-noarch).
- torch-ext/relu_tpu/__init__.py registers a plain `relu` callable via
  jax_op("relu_tpu::relu", _jax_relu) and re-exports the layers submodule.
- torch-ext/relu_tpu/layers/__init__.py exposes ReLU as a torch.nn.Module
  wrapping `relu` so kernelize(device="tpu") can swap a forward in.
- CARD.md mirrors the templated CARD.md committed under examples so
  kernel-builder can publish without further edits.
---
 examples/kernels/relu-tpu/CARD.md             | 61 +++++++++++++++++++
 examples/kernels/relu-tpu/build.toml          | 16 +++++
 .../relu-tpu/torch-ext/relu_tpu/__init__.py   | 14 +++++
 .../torch-ext/relu_tpu/layers/__init__.py     |  9 +++
 4 files changed, 100 insertions(+)
 create mode 100644 examples/kernels/relu-tpu/CARD.md
 create mode 100644 examples/kernels/relu-tpu/build.toml
 create mode 100644 examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py
 create mode 100644 examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py

diff --git a/examples/kernels/relu-tpu/CARD.md b/examples/kernels/relu-tpu/CARD.md
new file mode 100644
index 00000000..92ecc33d
--- /dev/null
+++ b/examples/kernels/relu-tpu/CARD.md
@@ -0,0 +1,61 @@
+---
+library_name: kernels
+{% if license %}license: {{ license }}
+{% endif %}---
+
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+
+## How to use
+{% if functions %}
+
+```python
+# make sure `kernels` is installed: `pip install -U kernels`
+from kernels import get_kernel
+
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
+
+{{ functions[0] }}(...)
+```
+{% else %}
+
+Usage example not available.
+{% endif %}
+
+## Available functions
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
+
+## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
+
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}
diff --git a/examples/kernels/relu-tpu/build.toml b/examples/kernels/relu-tpu/build.toml
new file mode 100644
index 00000000..c930a6a6
--- /dev/null
+++ b/examples/kernels/relu-tpu/build.toml
@@ -0,0 +1,16 @@
+[general]
+name = "relu-tpu"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = ["tpu"]
+
+[general.hub]
+repo-id = "kernels-test/relu-tpu"
+
+[general.tpu]
+python-depends = ["torch_tpu", "jax"]
+
+[torch-noarch]
+
+[kernel]
diff --git a/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py b/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py
new file mode 100644
index 00000000..e4d7224c
--- /dev/null
+++ b/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py
@@ -0,0 +1,14 @@
+import jax
+import jax.numpy as jnp
+from torch_tpu._internal.pallas import jax_op
+
+
+def _jax_relu(x):
+    return jnp.maximum(x, 0)
+
+
+relu = jax_op("relu_tpu::relu", _jax_relu)
+
+from . import layers  # noqa: E402
+
+__all__ = ["layers", "relu"]
diff --git a/examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py b/examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py
new file mode 100644
index 00000000..7800a2ef
--- /dev/null
+++ b/examples/kernels/relu-tpu/torch-ext/relu_tpu/layers/__init__.py
@@ -0,0 +1,9 @@
+import torch
+import torch.nn as nn
+
+from .. import relu
+
+
+class ReLU(nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return relu(x)

From 36a9d72cdf06b4abd5d0e8e6fdd6bab082f41547 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 6 Jul 2026 15:16:14 +0000
Subject: [PATCH 05/13] kernels/pytest: register tpu_only marker

test_basic.py and test_layer.py now use @pytest.mark.tpu_only to gate
tests that require a torch_tpu interpreter. Register the marker so
pytest will not warn about an unknown marker when those tests run on a
non-TPU host and the tpu_only skip is exercised.
---
 kernels/pytest.ini | 1 +
 1 file changed, 1 insertion(+)

diff --git a/kernels/pytest.ini b/kernels/pytest.ini
index a9421163..22164adc 100644
--- a/kernels/pytest.ini
+++ b/kernels/pytest.ini
@@ -5,4 +5,5 @@ markers =
     darwin_only: marks tests that should only run on macOS
     xpu_only: marks tests that should only run on hosts with Intel XPUs
     npu_only: marks tests that should only run on Ascend NPUs
+    tpu_only: marks tests that should only run on hosts with TPU devices
     token: enable tests that require a write token

From 43596557094417d9ccbbd53c1bd9231e81c81cb3 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 7 Jul 2026 16:23:24 +0000
Subject: [PATCH 06/13] nix-builder: add TPU backend support

Add TPU as a first-class backend so a kernel repo declaring
`backends = ["tpu"]` produces a `torch-tpu` noarch variant, dev
shell, and CI/bundle outputs through the same generators as the
other five backends.

Plumbing:
- lib/kernel-config.nix, lib/torch-version-utils.nix: add `tpu` to
  the backend init attrsets and an `isTpu` predicate + `backend` arm.
- lib/mk-build-set.nix: add a `tpu` backendConfig and an empty
  backendOverlay arm (torch_tpu/jax/libtpu layer on the CPU torch
  wheel via pythonPackagesExtensions, not nixpkgs overlays). The tpu
  buildConfig sets allowUnfree = true because libtpu's wheel METADATA
  declares its license as "Google Cloud Platform Terms of Service"
  (unfree), unlike the Apache-2.0 torch_tpu.
- lib/gen-flake-outputs.nix: add a `tpu` arm to buildConfigBackend and
  onePerFramework; backendCi/backendBundle pick it up automatically.
- build-variants.json + scripts/gen_variants_markdown.py: register
  the `torch-tpu` noarch variant and the "TPU" platform name.
- overlay.nix: wire jaxlib, libtpu, torch_tpu into
  pythonPackagesExtensions.
- README.md: add a TPU row (Tier 3, experimental).

Nix packages for the TPU Python wheels:
- jax, jaxlib: PyPI wheels pinned to 0.10.2 (cp312).
- libtpu 0.0.43, torch_tpu 0.1.1.dev20260707090224: wheels fetched
  directly from Google's Artifact Registry (gcloud), which requires an
  OAuth2 bearer token passed at fetch time via GCLOUD_ACCESS_TOKEN and
  turned into a netrc entry. Hashes are pinned for the cp312 wheel;
  scripts/helpers/get_torch_tpu_hash.sh prefetches the sha256 for a
  given ABI tag when refreshing versions.
- lib/mk-build-set.nix: pin the noarch TPU extension env to python312
  (the wheels are cp312-only at this point; a later commit derives the
  ABI tag from the package set's python instead).
---
 docs/source/builder/build-variants.md         |  5 ++
 nix-builder/README.md                         |  3 +-
 nix-builder/build-variants.json               |  3 +
 nix-builder/lib/build-variants.nix            |  5 +-
 nix-builder/lib/gen-flake-outputs.nix         |  3 +
 nix-builder/lib/kernel-config.nix             |  2 +
 nix-builder/lib/mk-build-set.nix              | 15 ++++-
 nix-builder/lib/torch-version-utils.nix       |  3 +
 nix-builder/overlay.nix                       |  4 ++
 .../pkgs/python-modules/libtpu/default.nix    | 59 +++++++++++++++++++
 .../pkgs/python-modules/torch_tpu/default.nix | 59 +++++++++++++++++++
 nix-builder/scripts/gen_variants_markdown.py  |  1 +
 .../scripts/helpers/get_torch_tpu_hash.sh     | 31 ++++++++++
 13 files changed, 190 insertions(+), 3 deletions(-)
 create mode 100644 nix-builder/pkgs/python-modules/libtpu/default.nix
 create mode 100644 nix-builder/pkgs/python-modules/torch_tpu/default.nix
 create mode 100755 nix-builder/scripts/helpers/get_torch_tpu_hash.sh

diff --git a/docs/source/builder/build-variants.md b/docs/source/builder/build-variants.md
index bdc9033f..5c7011d7 100644
--- a/docs/source/builder/build-variants.md
+++ b/docs/source/builder/build-variants.md
@@ -62,6 +62,10 @@ available. This list will be updated as new PyTorch versions are released.
 - `torch213-cxx11-rocm71-x86_64-linux`
 - `torch213-cxx11-rocm72-x86_64-linux`
 
+## TPU x86_64-linux
+
+- `torch-tpu`
+
 ## XPU x86_64-linux
 
 - `torch211-cxx11-xpu20253-x86_64-linux`
@@ -77,4 +81,5 @@ one or more of the following variants:
 - `torch-cuda`
 - `torch-metal`
 - `torch-rocm`
+- `torch-tpu`
 - `torch-xpu`
diff --git a/nix-builder/README.md b/nix-builder/README.md
index 7b31003b..e59b7119 100644
--- a/nix-builder/README.md
+++ b/nix-builder/README.md
@@ -54,8 +54,9 @@ The compiled kernel will then be available in the local `build/` directory.
 | Metal      | ✓               | ✓                      | ✗                       | 2    |
 | Huawei NPU | ✓               | ✗                      | ✗                       | 3    |
 | Neuron     | ✓               | x                      | x                       | 3    |
+| TPU        | ✓               | ✗                      | ✗                       | 3    |
 
-**Warning:** Neuron support is experimental and currently requires pre-release packages.
+**Warning:** Neuron and TPU support are experimental and packages are not public yet.
 
 # 📚 Documentation
 
diff --git a/nix-builder/build-variants.json b/nix-builder/build-variants.json
index 7fa073eb..da668552 100644
--- a/nix-builder/build-variants.json
+++ b/nix-builder/build-variants.json
@@ -54,6 +54,9 @@
       "torch213-cxx11-rocm71-x86_64-linux",
       "torch213-cxx11-rocm72-x86_64-linux"
     ],
+    "tpu": [
+      "torch-tpu"
+    ],
     "xpu": [
       "torch211-cxx11-xpu20253-x86_64-linux",
       "torch212-cxx11-xpu20253-x86_64-linux",
diff --git a/nix-builder/lib/build-variants.nix b/nix-builder/lib/build-variants.nix
index f7dfb777..8948911f 100644
--- a/nix-builder/lib/build-variants.nix
+++ b/nix-builder/lib/build-variants.nix
@@ -26,7 +26,10 @@ rec {
           throw "No compute framework set in Torch version";
     in
     version:
-    if version.system == "aarch64-darwin" then
+    # TPU builds are noarch: one variant, independent of Torch version and system.
+    if backend version == "tpu" then
+      "torch-tpu"
+    else if version.system == "aarch64-darwin" then
       "torch${flattenVersion version.torchVersion}-${computeString version}-${version.system}"
     else
       "torch${flattenVersion version.torchVersion}-cxx11-${computeString version}-${version.system}";
diff --git a/nix-builder/lib/gen-flake-outputs.nix b/nix-builder/lib/gen-flake-outputs.nix
index 99fba216..9e553cb3 100644
--- a/nix-builder/lib/gen-flake-outputs.nix
+++ b/nix-builder/lib/gen-flake-outputs.nix
@@ -58,6 +58,8 @@ let
       "xpu"
     else if buildConfig.metal or false then
       "metal"
+    else if buildConfig.tpu or false then
+      "tpu"
     else
       throw "Cannot determine framework for build set";
 
@@ -249,6 +251,7 @@ in
             ++ (headOrEmpty (setsWithFramework "cuda"))
             ++ (headOrEmpty (setsWithFramework "metal"))
             ++ (headOrEmpty (setsWithFramework "rocm"))
+            ++ (headOrEmpty (setsWithFramework "tpu"))
             ++ (headOrEmpty (setsWithFramework "xpu"));
         in
         build.mkExtensionBundle {
diff --git a/nix-builder/lib/kernel-config.nix b/nix-builder/lib/kernel-config.nix
index 0b274135..4770404a 100644
--- a/nix-builder/lib/kernel-config.nix
+++ b/nix-builder/lib/kernel-config.nix
@@ -48,6 +48,7 @@ in
         cuda = false;
         metal = false;
         rocm = false;
+        tpu = false;
         xpu = false;
       };
     in
@@ -63,6 +64,7 @@ in
         cuda = false;
         metal = false;
         rocm = false;
+        tpu = false;
         xpu = false;
       };
     in
diff --git a/nix-builder/lib/mk-build-set.nix b/nix-builder/lib/mk-build-set.nix
index 920d51cc..e85eb92a 100644
--- a/nix-builder/lib/mk-build-set.nix
+++ b/nix-builder/lib/mk-build-set.nix
@@ -82,6 +82,14 @@ let
       rocmSupport = true;
     };
 
+    tpu = {
+      # torch_tpu is Apache-2.0, but libtpu's wheel METADATA declares
+      # its license as "Google Cloud Platform Terms of Service"
+      # (unfree), so the tpu buildSet needs allowUnfree just like the
+      # cuda/rocm/xpu sets. See pkgs/python-modules/libtpu/default.nix.
+      allowUnfree = true;
+    };
+
     xpu = {
       allowUnfree = true;
       xpuSupport = true;
@@ -117,6 +125,8 @@ let
       [ (overlayForCudaVersion cudaVersion ptxasVersion) ]
     else if buildConfig.backend == "rocm" then
       [ (overlayForRocmVersion rocmVersion) ]
+    else if buildConfig.backend == "tpu" then
+      [ ]
     else if buildConfig.backend == "metal" then
       [ ]
     else if buildConfig.backend == "xpu" then
@@ -138,7 +148,10 @@ let
 
   torch = pkgs.python3.pkgs.torch;
 
-  extension = pkgs.callPackage ./extension { inherit torch; };
+  extension = pkgs.callPackage ./extension {
+    inherit torch;
+    python3 = if buildConfig.backend == "tpu" then pkgs.python312 else pkgs.python3;
+  };
 
   variants = import ./variants {
     inherit lib buildConfig;
diff --git a/nix-builder/lib/torch-version-utils.nix b/nix-builder/lib/torch-version-utils.nix
index 5b43482c..b7009de5 100644
--- a/nix-builder/lib/torch-version-utils.nix
+++ b/nix-builder/lib/torch-version-utils.nix
@@ -4,6 +4,7 @@ let
   isCuda = version: version ? cudaVersion;
   isMetal = version: version.metal or false;
   isRocm = version: version ? rocmVersion;
+  isTpu = version: version.tpu or false;
   isXpu = version: version ? xpuVersion;
 
 in
@@ -25,6 +26,8 @@ rec {
       "metal"
     else if isRocm version then
       "rocm"
+    else if isTpu version then
+      "tpu"
     else if isXpu version then
       "xpu"
     else
diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix
index f9480651..34357e89 100644
--- a/nix-builder/overlay.nix
+++ b/nix-builder/overlay.nix
@@ -165,6 +165,10 @@ final: prev:
 
         jax-tvm-ffi = python-self.callPackage ./pkgs/python-modules/jax-tvm-ffi { };
 
+        libtpu = python-self.callPackage ./pkgs/python-modules/libtpu { };
+
+        torch_tpu = python-self.callPackage ./pkgs/python-modules/torch_tpu { };
+
         jupyter-server = python-super.jupyter-server.overrideAttrs (
           _: prevAttrs: {
             # Gets stuck sometimes, already tested in nixpkgs.
diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix
new file mode 100644
index 00000000..09109fe3
--- /dev/null
+++ b/nix-builder/pkgs/python-modules/libtpu/default.nix
@@ -0,0 +1,59 @@
+{ lib
+, buildPythonPackage
+, fetchurl
+, python312
+}:
+
+# Fetched directly from Google's Artifact Registry, which requires an
+# OAuth2 bearer token. The token is passed at fetch time through the
+# GCLOUD_ACCESS_TOKEN environment variable and turned into a netrc
+# entry, so it never appears in the URL or the .drv:
+#
+#   export GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token)
+#
+# (With a multi-user Nix daemon the variable must be visible to the
+# daemon, not the client shell.)
+#
+# libtpu ships the `libtpu/libtpu.so` runtime that both torch_tpu and
+# jaxlib dlopen. Pinned to 0.0.x to match the jax 0.10.x / torch_tpu
+# 0.1.x stack (torch_tpu requires libtpu>=0.0.40; jax's "tpu" extra
+# pins libtpu==0.0.43.*), so bump it together with jaxlib/torch_tpu.
+#
+# NOTE on license: libtpu's wheel METADATA declares its license as
+# "Google Cloud Platform Terms of Service" — i.e. unfree, NOT
+# Apache-2.0 (unlike torch_tpu). The tpu buildSet therefore sets
+# allowUnfree = true (see lib/mk-build-set.nix).
+
+buildPythonPackage rec {
+  pname = "libtpu";
+  version = "0.0.43";
+  format = "wheel";
+
+  src = fetchurl {
+    url = "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/libtpu/libtpu-${version}-cp312-cp312-manylinux_2_31_x86_64.whl";
+    hash = "sha256-cInxat7P+bjjy+vIVdBMU6bXEMXPSgdDggzCc9WNc0o=";
+    netrcImpureEnvVars = [ "GCLOUD_ACCESS_TOKEN" ];
+    netrcPhase = ''
+      if [ -z "''${GCLOUD_ACCESS_TOKEN:-}" ]; then
+        echo "GCLOUD_ACCESS_TOKEN is not set; cannot fetch libtpu." >&2
+        echo "Run: export GCLOUD_ACCESS_TOKEN=\$(gcloud auth print-access-token)" >&2
+        exit 1
+      fi
+      printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' \
+        "$GCLOUD_ACCESS_TOKEN" > netrc
+    '';
+  };
+
+  python = python312; # only cp312 wheels
+  dependencies = [ ];
+
+  pythonImportsCheck = [ "libtpu" ];
+  doInstallCheck = false; # requires actual /dev/accel
+
+  meta = with lib; {
+    description = "TPU runtime shared library, dlopened by torch_tpu and jaxlib";
+    homepage = "https://github.com/google-pytorch/torch_tpu";
+    license = licenses.unfree; # "Google Cloud Platform Terms of Service"
+    platforms = [ "x86_64-linux" ];
+  };
+}
diff --git a/nix-builder/pkgs/python-modules/torch_tpu/default.nix b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
new file mode 100644
index 00000000..24bbe46a
--- /dev/null
+++ b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
@@ -0,0 +1,59 @@
+{ lib
+, buildPythonPackage
+, fetchurl
+, python312
+
+# jax is not a hard dependency of the torch_tpu wheel, but TPU kernels
+# import torch_tpu._internal.pallas (jax_op), which needs jax at import
+# time, so it is listed here. jaxlib is not a torch_tpu dependency; it
+# is pulled in transitively by jax (nixpkgs' jax declares it).
+# libtpu is pinned in-tree via pkgs/python-modules/libtpu.
+, jax
+, libtpu
+}:
+
+# Fetched directly from Google's Artifact Registry, which requires an
+# OAuth2 bearer token. The token is passed at fetch time through the
+# GCLOUD_ACCESS_TOKEN environment variable and turned into a netrc
+# entry, so it never appears in the URL or the .drv:
+#
+#   export GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token)
+#
+# (With a multi-user Nix daemon the variable must be visible to the
+# daemon, not the client shell.) The dev build is dated; bump the date
+# suffix in `version` and re-run
+# scripts/helpers/get_torch_tpu_hash.sh when refreshing.
+
+buildPythonPackage rec {
+  pname = "torch_tpu";
+  version = "0.1.1.dev20260707090224";
+  format = "wheel";
+
+  src = fetchurl {
+    url = "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/torch-tpu/torch_tpu-${version}-cp312-cp312-manylinux_2_31_x86_64.whl";
+    hash = "sha256-aTgLY6w1Q6zZvGlBHppNmmUFVthtzD0zoV2nAUoOtBM=";
+    netrcImpureEnvVars = [ "GCLOUD_ACCESS_TOKEN" ];
+    netrcPhase = ''
+      if [ -z "''${GCLOUD_ACCESS_TOKEN:-}" ]; then
+        echo "GCLOUD_ACCESS_TOKEN is not set; cannot fetch torch_tpu." >&2
+        echo "Run: export GCLOUD_ACCESS_TOKEN=\$(gcloud auth print-access-token)" >&2
+        exit 1
+      fi
+      printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' \
+        "$GCLOUD_ACCESS_TOKEN" > netrc
+    '';
+  };
+
+  python = python312; # only cp312 wheels
+  dependencies = [ libtpu jax ];
+
+  pythonImportsCheck = [ "torch_tpu" ];
+  doInstallCheck = false; # requires actual /dev/accel
+
+  meta = with lib; {
+    description = "Torch TPU backend (PrivateUse1 name: \"tpu\")";
+    homepage = "https://github.com/google-pytorch/torch_tpu";
+    license = licenses.asl20;
+    platforms = [ "x86_64-linux" ];
+  };
+}
diff --git a/nix-builder/scripts/gen_variants_markdown.py b/nix-builder/scripts/gen_variants_markdown.py
index b097cebc..f528144c 100755
--- a/nix-builder/scripts/gen_variants_markdown.py
+++ b/nix-builder/scripts/gen_variants_markdown.py
@@ -8,6 +8,7 @@
     "cuda": "CUDA",
     "metal": "Metal",
     "rocm": "ROCm",
+    "tpu": "TPU",
     "xpu": "XPU",
 }
 
diff --git a/nix-builder/scripts/helpers/get_torch_tpu_hash.sh b/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
new file mode 100755
index 00000000..a8a65f5e
--- /dev/null
+++ b/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
@@ -0,0 +1,31 @@
+#!/usr/bin/env bash
+# Print the wheel URL and SRI hash for a torch-tpu-virtual-registry
+# package, so a developer can paste the result straight into
+# pkgs/python-modules/{torch_tpu,libtpu}/default.nix.
+#
+# The registry's simple index publishes each wheel's sha256 as a URL
+# fragment, so no download is needed to compute the hash.
+#
+# Usage:
+#   bash scripts/helpers/get_torch_tpu_hash.sh torch_tpu 0.1.1.dev20260707090224
+#   bash scripts/helpers/get_torch_tpu_hash.sh libtpu    0.0.43
+#
+# Requires `gcloud auth login` once, plus curl and python3.
+set -euo pipefail
+pkg="${1:?usage: $0 <pkg> <version>, e.g. $0 torch_tpu 0.1.1.dev20260707090224}"
+ver="${2:?usage: $0 <pkg> <version>, e.g. $0 torch_tpu 0.1.1.dev20260707090224}"
+token="$(gcloud auth print-access-token)"
+index="https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/simple/${pkg//_/-}/"
+index_html="$(curl -sSf -u "oauth2accesstoken:${token}" "${index}")"
+href="$(grep -o 'href="[^"]*"' <<<"${index_html}" \
+    | grep -F "/${pkg}-${ver}-cp312-" \
+    | sed 's/^href="//; s/"$//' \
+    | head -1 || true)"
+if [ -z "${href}" ]; then
+    echo "No cp312 wheel for ${pkg}==${ver} in ${index}" >&2
+    exit 1
+fi
+echo "URL:    ${href%%#*}"
+echo "Hash:   $(python3 -c \
+    "import base64,sys; print('sha256-' + base64.b64encode(bytes.fromhex(sys.argv[1])).decode())" \
+    "${href##*#sha256=}")"

From ef27628788afe0faa005fe4554a87cf81b66801f Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 16:38:41 +0000
Subject: [PATCH 07/13] nix-builder: add TPU build configuration to
 versions.nix

Without a versions.nix record carrying tpu = true no TPU build set is
ever instantiated, so kernels declaring backends = ["tpu"] resolved
to zero applicable build sets. Add one torch 2.11 entry (torch_tpu
pins torch>=2.11,<2.12), x86_64-linux only.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 nix-builder/versions.nix | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/nix-builder/versions.nix b/nix-builder/versions.nix
index 1618d26e..3eb87138 100644
--- a/nix-builder/versions.nix
+++ b/nix-builder/versions.nix
@@ -58,6 +58,12 @@
     systems = [ "x86_64-linux" ];
     bundleBuild = true;
   }
+  {
+    torchVersion = "2.11";
+    tpu = true;
+    systems = [ "x86_64-linux" ];
+    bundleBuild = true;
+  }
   {
     torchVersion = "2.11";
     xpuVersion = "2025.3.2";

From 38712a18454ef596bf539befd4fcbe7cbe6c3489 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 16:38:41 +0000
Subject: [PATCH 08/13] examples: add flake.nix to relu-tpu kernel

Mirrors the standard example flake (relu) so the kernel can be built
with nix build directly, producing backendBundle.tpu and the
redistributable torch-tpu variant.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 examples/kernels/relu-tpu/flake.nix | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)
 create mode 100644 examples/kernels/relu-tpu/flake.nix

diff --git a/examples/kernels/relu-tpu/flake.nix b/examples/kernels/relu-tpu/flake.nix
new file mode 100644
index 00000000..a3087256
--- /dev/null
+++ b/examples/kernels/relu-tpu/flake.nix
@@ -0,0 +1,17 @@
+{
+  description = "Flake for TPU ReLU kernel";
+
+  inputs = {
+    kernel-builder.url = "path:../../..";
+  };
+
+  outputs =
+    {
+      self,
+      kernel-builder,
+    }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

From ee10f6977c102dcafb82bc8e1c236d799bc4abd1 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 16:39:00 +0000
Subject: [PATCH 09/13] nix-builder: build TPU wheels against the default
 python

The torch binary rewrap ships one wheel per (version, system,
framework) keyed to the nixpkgs default python (currently 3.13), so
pinning the TPU extension env to python312 forced an unusable cp313
torch into a 3.12 env. The registry publishes torch_tpu and libtpu
wheels for cp311..cp314, so drop the python312 pin and derive the
wheel ABI tag from the package set's python instead (hashes pinned
for cp313).

Also:
- accept the tpu attribute in mk-build-set's buildConfig signature
  (build configs from versions.nix carry it)
- declare torch_tpu's full wheel Requires-Dist set (torch, portpicker,
  tensorboard, frozendict, immutabledict; numpy/absl-py come via jax)
- autoPatchelfHook with torch's lib dir on the search path, since the
  bundled extensions link libtorch_python.so/libc10.so/libtorch_cpu.so
- get_torch_tpu_hash.sh takes the ABI tag as an optional argument

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 nix-builder/lib/mk-build-set.nix              |  6 +-
 .../pkgs/python-modules/libtpu/default.nix    | 23 ++++---
 .../pkgs/python-modules/torch_tpu/default.nix | 66 ++++++++++++++-----
 .../scripts/helpers/get_torch_tpu_hash.sh     | 14 ++--
 4 files changed, 76 insertions(+), 33 deletions(-)

diff --git a/nix-builder/lib/mk-build-set.nix b/nix-builder/lib/mk-build-set.nix
index e85eb92a..2b0355c6 100644
--- a/nix-builder/lib/mk-build-set.nix
+++ b/nix-builder/lib/mk-build-set.nix
@@ -110,6 +110,7 @@ buildConfig@{
   ptxasVersion ? cudaVersion,
   metal ? false,
   rocmVersion ? null,
+  tpu ? false,
   xpuVersion ? null,
   torchVersion,
   system,
@@ -148,10 +149,7 @@ let
 
   torch = pkgs.python3.pkgs.torch;
 
-  extension = pkgs.callPackage ./extension {
-    inherit torch;
-    python3 = if buildConfig.backend == "tpu" then pkgs.python312 else pkgs.python3;
-  };
+  extension = pkgs.callPackage ./extension { inherit torch; };
 
   variants = import ./variants {
     inherit lib buildConfig;
diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix
index 09109fe3..53cc24f0 100644
--- a/nix-builder/pkgs/python-modules/libtpu/default.nix
+++ b/nix-builder/pkgs/python-modules/libtpu/default.nix
@@ -1,7 +1,8 @@
-{ lib
-, buildPythonPackage
-, fetchurl
-, python312
+{
+  lib,
+  buildPythonPackage,
+  fetchurl,
+  python,
 }:
 
 # Fetched directly from Google's Artifact Registry, which requires an
@@ -24,14 +25,21 @@
 # Apache-2.0 (unlike torch_tpu). The tpu buildSet therefore sets
 # allowUnfree = true (see lib/mk-build-set.nix).
 
+let
+  # The wheel ships per-CPython-ABI builds (cp311..cp314); pick the tag
+  # matching the python this package set is built for. The hash below is
+  # for cp313 (the nixpkgs default python); re-run
+  # scripts/helpers/get_torch_tpu_hash.sh when either moves.
+  abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
+in
 buildPythonPackage rec {
   pname = "libtpu";
   version = "0.0.43";
   format = "wheel";
 
   src = fetchurl {
-    url = "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/libtpu/libtpu-${version}-cp312-cp312-manylinux_2_31_x86_64.whl";
-    hash = "sha256-cInxat7P+bjjy+vIVdBMU6bXEMXPSgdDggzCc9WNc0o=";
+    url = "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/libtpu/libtpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
+    hash = "sha256-X5LVmwJuRcNkMG+m7kKgcaOSKnWjeuw4HtCU+lfWamY="; # cp313
     netrcImpureEnvVars = [ "GCLOUD_ACCESS_TOKEN" ];
     netrcPhase = ''
       if [ -z "''${GCLOUD_ACCESS_TOKEN:-}" ]; then
@@ -44,7 +52,6 @@ buildPythonPackage rec {
     '';
   };
 
-  python = python312; # only cp312 wheels
   dependencies = [ ];
 
   pythonImportsCheck = [ "libtpu" ];
@@ -52,7 +59,7 @@ buildPythonPackage rec {
 
   meta = with lib; {
     description = "TPU runtime shared library, dlopened by torch_tpu and jaxlib";
-    homepage = "https://github.com/google-pytorch/torch_tpu";
+    homepage = "https://cloud.google.com/tpu";
     license = licenses.unfree; # "Google Cloud Platform Terms of Service"
     platforms = [ "x86_64-linux" ];
   };
diff --git a/nix-builder/pkgs/python-modules/torch_tpu/default.nix b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
index 24bbe46a..caa24ec6 100644
--- a/nix-builder/pkgs/python-modules/torch_tpu/default.nix
+++ b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
@@ -1,15 +1,26 @@
-{ lib
-, buildPythonPackage
-, fetchurl
-, python312
-
-# jax is not a hard dependency of the torch_tpu wheel, but TPU kernels
-# import torch_tpu._internal.pallas (jax_op), which needs jax at import
-# time, so it is listed here. jaxlib is not a torch_tpu dependency; it
-# is pulled in transitively by jax (nixpkgs' jax declares it).
-# libtpu is pinned in-tree via pkgs/python-modules/libtpu.
-, jax
-, libtpu
+{
+  lib,
+  autoPatchelfHook,
+  buildPythonPackage,
+  fetchurl,
+  python,
+  stdenv,
+
+  # jax is not a hard dependency of the torch_tpu wheel, but TPU kernels
+  # import torch_tpu._internal.pallas (jax_op), which needs jax at import
+  # time, so it is listed here. jaxlib is not a torch_tpu dependency; it
+  # is pulled in transitively by jax (nixpkgs' jax declares it).
+  # libtpu is pinned in-tree via pkgs/python-modules/libtpu.
+  jax,
+  libtpu,
+
+  # The remaining wheel Requires-Dist entries (numpy and absl-py come in
+  # transitively through jax).
+  frozendict,
+  immutabledict,
+  portpicker,
+  tensorboard,
+  torch,
 }:
 
 # Fetched directly from Google's Artifact Registry, which requires an
@@ -24,14 +35,21 @@
 # suffix in `version` and re-run
 # scripts/helpers/get_torch_tpu_hash.sh when refreshing.
 
+let
+  # The wheel ships per-CPython-ABI builds (cp311..cp314); pick the tag
+  # matching the python this package set is built for. The hash below is
+  # for cp313 (the nixpkgs default python); re-run
+  # scripts/helpers/get_torch_tpu_hash.sh when either moves.
+  abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
+in
 buildPythonPackage rec {
   pname = "torch_tpu";
   version = "0.1.1.dev20260707090224";
   format = "wheel";
 
   src = fetchurl {
-    url = "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/torch-tpu/torch_tpu-${version}-cp312-cp312-manylinux_2_31_x86_64.whl";
-    hash = "sha256-aTgLY6w1Q6zZvGlBHppNmmUFVthtzD0zoV2nAUoOtBM=";
+    url = "https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/torch-tpu/torch_tpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
+    hash = "sha256-eCjwoX0UKd/L/IccxXn88p0GdyQjFdPl2Er1luZz4H0="; # cp313
     netrcImpureEnvVars = [ "GCLOUD_ACCESS_TOKEN" ];
     netrcPhase = ''
       if [ -z "''${GCLOUD_ACCESS_TOKEN:-}" ]; then
@@ -44,8 +62,24 @@ buildPythonPackage rec {
     '';
   };
 
-  python = python312; # only cp312 wheels
-  dependencies = [ libtpu jax ];
+  dependencies = [
+    libtpu
+    jax
+    frozendict
+    immutabledict
+    portpicker
+    tensorboard
+    torch
+  ];
+
+  # The bundled extensions link against libtorch_python.so / libc10.so /
+  # libtorch_cpu.so, which live in torch's site-packages, not on the
+  # default search path.
+  nativeBuildInputs = [ autoPatchelfHook ];
+  buildInputs = [ stdenv.cc.cc.lib ];
+  preFixup = ''
+    addAutoPatchelfSearchPath ${torch}/${python.sitePackages}/torch/lib
+  '';
 
   pythonImportsCheck = [ "torch_tpu" ];
   doInstallCheck = false; # requires actual /dev/accel
diff --git a/nix-builder/scripts/helpers/get_torch_tpu_hash.sh b/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
index a8a65f5e..ee3f9686 100755
--- a/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
+++ b/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
@@ -8,21 +8,25 @@
 #
 # Usage:
 #   bash scripts/helpers/get_torch_tpu_hash.sh torch_tpu 0.1.1.dev20260707090224
-#   bash scripts/helpers/get_torch_tpu_hash.sh libtpu    0.0.43
+#   bash scripts/helpers/get_torch_tpu_hash.sh libtpu    0.0.43 cp313
+#
+# The third argument is the CPython ABI tag (default cp313, matching the
+# nixpkgs default python used by the builder).
 #
 # Requires `gcloud auth login` once, plus curl and python3.
 set -euo pipefail
-pkg="${1:?usage: $0 <pkg> <version>, e.g. $0 torch_tpu 0.1.1.dev20260707090224}"
-ver="${2:?usage: $0 <pkg> <version>, e.g. $0 torch_tpu 0.1.1.dev20260707090224}"
+pkg="${1:?usage: $0 <pkg> <version> [abi], e.g. $0 torch_tpu 0.1.1.dev20260707090224}"
+ver="${2:?usage: $0 <pkg> <version> [abi], e.g. $0 torch_tpu 0.1.1.dev20260707090224}"
+abi="${3:-cp313}"
 token="$(gcloud auth print-access-token)"
 index="https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/simple/${pkg//_/-}/"
 index_html="$(curl -sSf -u "oauth2accesstoken:${token}" "${index}")"
 href="$(grep -o 'href="[^"]*"' <<<"${index_html}" \
-    | grep -F "/${pkg}-${ver}-cp312-" \
+    | grep -F "/${pkg}-${ver}-${abi}-" \
     | sed 's/^href="//; s/"$//' \
     | head -1 || true)"
 if [ -z "${href}" ]; then
-    echo "No cp312 wheel for ${pkg}==${ver} in ${index}" >&2
+    echo "No ${abi} wheel for ${pkg}==${ver} in ${index}" >&2
     exit 1
 fi
 echo "URL:    ${href%%#*}"

From 3cd1ac2aa1d2543299f140bc7bd29528cf6d3dc5 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 16:39:00 +0000
Subject: [PATCH 10/13] kernels-data,kernels: resolve TPU python deps to nix
 packages

torch_tpu and jax had empty nix package lists in the dependency maps,
so nix-builder's dependency resolution produced a check environment
without them and the sandboxed get_kernel check could not detect the
TPU backend. Point them at the torch_tpu/jax packages provided by the
nix-builder overlay.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 kernels-data/src/python_dependencies.json | 4 ++--
 kernels/src/kernels/python_depends.json   | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json
index a37bf8d6..8de880f5 100644
--- a/kernels-data/src/python_dependencies.json
+++ b/kernels-data/src/python_dependencies.json
@@ -37,11 +37,11 @@
     },
     "tpu": {
       "torch_tpu": {
-        "nix": [],
+        "nix": ["torch_tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
       },
       "jax": {
-        "nix": [],
+        "nix": ["jax"],
         "python": [{ "pkg": "jax", "import": "jax" }]
       },
       "libtpu": {
diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json
index a37bf8d6..8de880f5 100644
--- a/kernels/src/kernels/python_depends.json
+++ b/kernels/src/kernels/python_depends.json
@@ -37,11 +37,11 @@
     },
     "tpu": {
       "torch_tpu": {
-        "nix": [],
+        "nix": ["torch_tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
       },
       "jax": {
-        "nix": [],
+        "nix": ["jax"],
         "python": [{ "pkg": "jax", "import": "jax" }]
       },
       "libtpu": {

From cae1a7d3ef786507d06c8f955f6f911da408403f Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 16:50:25 +0000
Subject: [PATCH 11/13] kernels,kernel-builder: detect TPU via
 torch.backends.tpu
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

torch_tpu only registers the torch.tpu namespace when TPU hardware is
present (its loader is gated on the device count), so probing
hasattr(torch, "tpu") made backend detection hardware-dependent —
unlike every other backend, where detection reflects the installed
torch stack (e.g. torch.version.cuda on GPU-less hosts). This broke
nix-builder's sandboxed get_kernel check, which runs with torch_tpu
installed but no TPU device.

torch.backends.tpu is set unconditionally when torch_tpu is imported
(torch's device-backend autoload triggers this on import torch), so
probe that instead. Tests that need real hardware keep using the
device-count-based has_tpu fixture.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 .../pyproject/templates/torch/noarch/_ops.py  |  7 ++-
 kernels/src/kernels/backends.py               | 44 ++++++++++---------
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py
index 46c697e9..e8d30b5f 100644
--- a/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py
+++ b/kernel-builder/src/pyproject/templates/torch/noarch/_ops.py
@@ -4,8 +4,11 @@ def get_backend() -> str:
     """Detect the backend by inspecting torch."""
     import torch
 
-    if hasattr(torch, "tpu"):
-        # torch_tpu registers a torch.tpu namespace on import.
+    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.
         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 3aa6fbec..c5b82920 100644
--- a/kernels/src/kernels/backends.py
+++ b/kernels/src/kernels/backends.py
@@ -127,24 +127,6 @@ def parse(s: str) -> "Neuron":
         return Neuron()
 
 
-@strict
-@dataclass(unsafe_hash=True)
-class TPU:
-    @property
-    def name(self) -> str:
-        return "tpu"
-
-    @property
-    def variant_str(self) -> str:
-        return "tpu"
-
-    @staticmethod
-    def parse(s: str) -> "TPU":
-        if s != "tpu":
-            raise ValueError(f"Invalid TPU variant string: {s!r}")
-        return TPU()
-
-
 @dataclass(unsafe_hash=True)
 class ROCm:
     _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"rocm(\d+)(\d+)")
@@ -167,6 +149,24 @@ def parse(s: str) -> "ROCm":
         return ROCm(version=Version(f"{m.group(1)}.{m.group(2)}"))
 
 
+@strict
+@dataclass(unsafe_hash=True)
+class TPU:
+    @property
+    def name(self) -> str:
+        return "tpu"
+
+    @property
+    def variant_str(self) -> str:
+        return "tpu"
+
+    @staticmethod
+    def parse(s: str) -> "TPU":
+        if s != "tpu":
+            raise ValueError(f"Invalid TPU variant string: {s!r}")
+        return TPU()
+
+
 @dataclass(unsafe_hash=True)
 class XPU:
     _VARIANT_REGEX: ClassVar[re.Pattern] = re.compile(r"xpu(\d+)(\d+)")
@@ -215,8 +215,12 @@ def _backend() -> Backend:
     if has_torch:
         import torch
 
-        if hasattr(torch, "tpu"):
-            # torch_tpu registers a torch.tpu namespace on import.
+        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.
             return TPU()
         elif hasattr(torch, "neuron"):
             # Needs to be sorted before specific Torch builds, since Neuron

From 4bf394620a2d8a5fa0ecbf65fade852b8db19cab Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 17:02:26 +0000
Subject: [PATCH 12/13] kernels-data: include TPU python-depends in backend
 metadata

backend_python_depends only matched Cuda and Xpu, so a build.toml
[general.tpu] python-depends section never reached the generated
metadata.json and installed TPU kernels skipped dependency
validation. (Neuron has the same pre-existing gap; left untouched.)

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 kernels-data/src/config/mod.rs | 52 ++++++++++++++++++++++++++++++++++
 1 file changed, 52 insertions(+)

diff --git a/kernels-data/src/config/mod.rs b/kernels-data/src/config/mod.rs
index 3300b48c..ed111d70 100644
--- a/kernels-data/src/config/mod.rs
+++ b/kernels-data/src/config/mod.rs
@@ -145,6 +145,10 @@ impl General {
                 .cuda
                 .as_ref()
                 .and_then(|cuda| cuda.python_depends.as_ref()),
+            Backend::Tpu => self
+                .tpu
+                .as_ref()
+                .and_then(|tpu| tpu.python_depends.as_ref()),
             Backend::Xpu => self
                 .xpu
                 .as_ref()
@@ -454,3 +458,51 @@ pub enum ConfigError {
     #[error("Cannot migrate configuration: {reason:?}")]
     Migration { reason: String },
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn general_with_tpu_deps() -> General {
+        General {
+            name: KernelName::new("test-kernel").unwrap(),
+            version: 1,
+            license: "apache-2.0".to_string(),
+            upstream: None,
+            source: None,
+            backends: vec![Backend::Tpu],
+            hub: None,
+            python_depends: None,
+            cuda: None,
+            neuron: None,
+            tpu: Some(TpuGeneral {
+                python_depends: Some(vec!["torch_tpu".to_string()]),
+            }),
+            xpu: None,
+        }
+    }
+
+    #[test]
+    fn backend_python_depends_resolves_tpu() {
+        let general = general_with_tpu_deps();
+        let deps = general
+            .backend_python_depends(Backend::Tpu)
+            .map(|dep| dep.map(|(name, _)| name.to_string()))
+            .collect::<Result<Vec<_>>>()
+            .unwrap();
+
+        assert_eq!(deps, vec!["torch_tpu".to_string()]);
+    }
+
+    #[test]
+    fn backend_python_depends_empty_for_backend_without_deps() {
+        let general = general_with_tpu_deps();
+
+        assert!(
+            general
+                .backend_python_depends(Backend::Cpu)
+                .next()
+                .is_none()
+        );
+    }
+}

From 7126d26f2cdebfc7b1dd56d91f6a75dbd56015b9 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Fri, 10 Jul 2026 16:53:39 +0000
Subject: [PATCH 13/13] examples/relu-tpu: annotate the jax function signature

torch_tpu's jax_op verifies the wrapped function's signature and
rejects unannotated arguments, so the kernel failed to import.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py b/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py
index e4d7224c..677bca5f 100644
--- a/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py
+++ b/examples/kernels/relu-tpu/torch-ext/relu_tpu/__init__.py
@@ -3,7 +3,7 @@
 from torch_tpu._internal.pallas import jax_op
 
 
-def _jax_relu(x):
+def _jax_relu(x: jax.Array) -> jax.Array:
     return jnp.maximum(x, 0)