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/32] 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 266904e82..a37bf8d6a 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 5dff532a1..3aa6fbec8 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 5294a5f5d..60d221d6a 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 fd02707c2..c53c833bb 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 266904e82..a37bf8d6a 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 a9cb3c01c..513261ffd 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 1460f9454..5c9ce8dbd 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 e67abfc28..790e64458 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 008e2c8a2..1707c720e 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/32] 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 eb4324819..2c4dd34f9 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 20e29619f..43b7af6a6 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 9aa29f06d..3cfcae33e 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 29bb52fe5..87f9d1119 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 bc7855f71..3300b48c5 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 eed543ba6..f79204a8b 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 590d16e6a..766a0cf9f 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 a88db69c9..ecc6c0125 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 0154eeef6..3e8a99960 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 d895f7fbb..3da04dc46 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/32] 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 3fd7d7644..532bc8ba1 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 8ce37f48a..f12b4edd6 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 75fd8e086..46c697e9a 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/32] 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 000000000..92ecc33da
--- /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 000000000..c930a6a61
--- /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 000000000..e4d7224c5
--- /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 000000000..7800a2ef7
--- /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/32] 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 a94211637..22164adcc 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/32] 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 bdc9033fb..5c7011d76 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 7b31003ba..e59b7119e 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 7fa073eba..da668552c 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 f7dfb7776..8948911f1 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 99fba216e..9e553cb3d 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 0b274135f..4770404a6 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 920d51cc1..e85eb92a5 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 5b43482cb..b7009de5f 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 f94806512..34357e892 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 000000000..09109fe3e
--- /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 000000000..24bbe46aa
--- /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 b097cebcc..f528144ce 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 000000000..a8a65f5e5
--- /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/32] 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 1618d26e1..3eb871387 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/32] 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 000000000..a30872561
--- /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/32] 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 e85eb92a5..2b0355c63 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 09109fe3e..53cc24f01 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 24bbe46aa..caa24ec62 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 a8a65f5e5..ee3f96865 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/32] 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 a37bf8d6a..8de880f5b 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 a37bf8d6a..8de880f5b 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/32] 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 46c697e9a..e8d30b5ff 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 3aa6fbec8..c5b829209 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/32] 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 3300b48c5..ed111d70d 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/32] 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 e4d7224c5..677bca5ff 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)
 
 

From 68fe98d945195c46559e8c7836d0a62a120b8fa1 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 11:45:21 +0200
Subject: [PATCH 14/32] nix-builder: drop get_torch_tpu_hash.sh helper script

Inline the hash-refresh command as a comment in each default.nix
instead, matching the nix-prefetch-url precedent already used for
torch's binary wheels. `nix store prefetch-file --json | jq -r
.hash` emits the SRI hash directly, so there's no need for the
script's hand-rolled hex-to-base64 step.
---
 .../pkgs/python-modules/libtpu/default.nix    | 12 +++++--
 .../pkgs/python-modules/torch_tpu/default.nix | 12 ++++---
 .../scripts/helpers/get_torch_tpu_hash.sh     | 35 -------------------
 3 files changed, 17 insertions(+), 42 deletions(-)
 delete mode 100755 nix-builder/scripts/helpers/get_torch_tpu_hash.sh

diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix
index 53cc24f01..3cc57270c 100644
--- a/nix-builder/pkgs/python-modules/libtpu/default.nix
+++ b/nix-builder/pkgs/python-modules/libtpu/default.nix
@@ -13,7 +13,13 @@
 #   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.)
+# daemon, not the client shell.) To bump the version, update `version`
+# below, then recompute `hash` with:
+#
+#   nix store prefetch-file --json --hash-type sha256 \
+#     --option netrc-file <(printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' "$(gcloud auth print-access-token)") \
+#     'https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/libtpu/libtpu-<version>-<abi>-<abi>-manylinux_2_31_x86_64.whl' \
+#     | jq -r .hash
 #
 # 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
@@ -28,8 +34,8 @@
 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.
+  # for cp313 (the nixpkgs default python); if either moves, recompute
+  # via the command above.
   abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
 in
 buildPythonPackage rec {
diff --git a/nix-builder/pkgs/python-modules/torch_tpu/default.nix b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
index caa24ec62..5cf8f2e4d 100644
--- a/nix-builder/pkgs/python-modules/torch_tpu/default.nix
+++ b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
@@ -32,14 +32,18 @@
 #
 # (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.
+# suffix in `version` below, then recompute `hash` with:
+#
+#   nix store prefetch-file --json --hash-type sha256 \
+#     --option netrc-file <(printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' "$(gcloud auth print-access-token)") \
+#     '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' \
+#     | jq -r .hash
 
 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.
+  # for cp313 (the nixpkgs default python); if either moves, recompute
+  # via the command above.
   abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
 in
 buildPythonPackage rec {
diff --git a/nix-builder/scripts/helpers/get_torch_tpu_hash.sh b/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
deleted file mode 100755
index ee3f96865..000000000
--- a/nix-builder/scripts/helpers/get_torch_tpu_hash.sh
+++ /dev/null
@@ -1,35 +0,0 @@
-#!/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 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> [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}-${abi}-" \
-    | sed 's/^href="//; s/"$//' \
-    | head -1 || true)"
-if [ -z "${href}" ]; then
-    echo "No ${abi} 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 ad707356383ad3fbe65b14173898166a87c0876d Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 12:05:28 +0200
Subject: [PATCH 15/32] kernel-builder: drop tpu from init --backends help text

kernel_init's build_init_fileset has no has_tpu handling, so
`init --backends tpu` would silently emit the wrong scaffold (C++
torch-ext binding stubs instead of the torch-noarch/[general.tpu]
layout TPU repos actually need, as seen in relu-tpu). Don't
advertise it until init actually supports scaffolding it.
---
 kernel-builder/src/init.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kernel-builder/src/init.rs b/kernel-builder/src/init.rs
index f12b4edd6..8ce37f48a 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`, `tpu`, `xpu`).
+    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`).
     #[arg(long, num_args = 1.., default_values_t = default_init_backends())]
     pub backends: Vec<BackendSelection>,
 

From 226e1e4b6494a5decdefa7f2202c3feaab4bec81 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 12:05:40 +0200
Subject: [PATCH 16/32] kernel-builder: drop neuron from init --backends help
 text

Same gap as tpu: build_init_fileset has no has_neuron handling or
neuron template, so the generated scaffold doesn't match what a
real neuron repo needs (relu-nki uses [general.neuron] +
[torch-noarch], not the C++ torch-ext binding stubs init renders
today). cann is already left out of this list for the same reason;
neuron should be too.
---
 kernel-builder/src/init.rs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/kernel-builder/src/init.rs b/kernel-builder/src/init.rs
index 8ce37f48a..393de2dd9 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`, `rocm`, `xpu`).
     #[arg(long, num_args = 1.., default_values_t = default_init_backends())]
     pub backends: Vec<BackendSelection>,
 

From 6f1e3a54fbdefcb5ab199359688aac8a8faf460d Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 12:42:45 +0200
Subject: [PATCH 17/32] docs: regenerate builder-cli.md for init --backends
 help text

Fixes CI's "Check docs are up to date" job, which went stale after
tpu and neuron were dropped from the --backends help text.
---
 docs/source/builder-cli.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/builder-cli.md b/docs/source/builder-cli.md
index a6b8c726f..63c7a6ef7 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`, `tpu`, `xpu`)
+* `--backends <BACKENDS>` — Backends to enable (`all`, `cpu`, `cuda`, `metal`, `rocm`, `xpu`)
 
 * `--overwrite` — Overwrite existing scaffold files (preserves other files)
 

From fb05cac62b5156480c23992bae231ee43df33c84 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 14:18:01 +0000
Subject: [PATCH 18/32] nix-builder: exclude tpu from the shared build cache

torch_tpu's autoPatchelfHook step produces an output that is itself a
modified copy of Google's gated libtpu/torch_tpu wheels, so it must
not be published to the shared cache until public builds exist.

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

diff --git a/nix-builder/versions.nix b/nix-builder/versions.nix
index 3eb871387..316716d2a 100644
--- a/nix-builder/versions.nix
+++ b/nix-builder/versions.nix
@@ -62,7 +62,10 @@
     torchVersion = "2.11";
     tpu = true;
     systems = [ "x86_64-linux" ];
-    bundleBuild = true;
+    # Excluded from the shared build cache: torch_tpu's autoPatchelfHook
+    # step produces an output that is itself a modified copy of Google's
+    # gated libtpu/torch_tpu wheels, so it can't be published publicly.
+    bundleBuild = false;
   }
   {
     torchVersion = "2.11";

From da257e5bbae1587f53d840616c9e7dd688afb381 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 14:18:12 +0000
Subject: [PATCH 19/32] nix-builder: fetch libtpu/torch_tpu via requireFile,
 not fetchurl

The gated Google Artifact Registry fetch previously ran inside the
build via fetchurl's netrcImpureEnvVars, requiring an impure build.
Switch to requireFile plus a prefetch-tpu-wheel.sh helper so the
authenticated fetch happens out-of-band and the Nix build itself
stays pure.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 .../pkgs/python-modules/libtpu/default.nix    | 46 ++++++++-----------
 .../pkgs/python-modules/torch_tpu/default.nix | 46 ++++++++-----------
 .../scripts/helpers/prefetch-tpu-wheel.sh     | 40 ++++++++++++++++
 3 files changed, 80 insertions(+), 52 deletions(-)
 create mode 100755 nix-builder/scripts/helpers/prefetch-tpu-wheel.sh

diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix
index 3cc57270c..ff18bdbab 100644
--- a/nix-builder/pkgs/python-modules/libtpu/default.nix
+++ b/nix-builder/pkgs/python-modules/libtpu/default.nix
@@ -1,25 +1,22 @@
 {
   lib,
   buildPythonPackage,
-  fetchurl,
+  requireFile,
   python,
 }:
 
-# 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:
+# libtpu is served from Google's gated Artifact Registry, which requires
+# an OAuth2 bearer token. That token can't be handed to a build sandbox
+# without making the derivation impure, so instead the wheel is
+# prefetched out-of-band and handed to `requireFile`, which only checks
+# that a file matching `name` + `hash` already exists in the store:
 #
-#   export GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token)
+#   GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+#     nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43
 #
-# (With a multi-user Nix daemon the variable must be visible to the
-# daemon, not the client shell.) To bump the version, update `version`
-# below, then recompute `hash` with:
-#
-#   nix store prefetch-file --json --hash-type sha256 \
-#     --option netrc-file <(printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' "$(gcloud auth print-access-token)") \
-#     'https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/libtpu/libtpu-<version>-<abi>-<abi>-manylinux_2_31_x86_64.whl' \
-#     | jq -r .hash
+# The script prints the `hash` line to paste in below. To bump the
+# version, update `version`, then re-run the script with the new
+# version.
 #
 # 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
@@ -35,7 +32,7 @@ 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); if either moves, recompute
-  # via the command above.
+  # via the prefetch script above.
   abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
 in
 buildPythonPackage rec {
@@ -43,18 +40,15 @@ buildPythonPackage rec {
   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}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
+  src = requireFile {
+    name = "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
-        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
+    message = ''
+      libtpu is served from a gated Google Artifact Registry and cannot
+      be fetched by a pure Nix build. Fetch and register it with:
+
+        GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+          nix-builder/scripts/helpers/prefetch-tpu-wheel.sh libtpu ${version} ${abi}
     '';
   };
 
diff --git a/nix-builder/pkgs/python-modules/torch_tpu/default.nix b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
index 5cf8f2e4d..c5dc007f9 100644
--- a/nix-builder/pkgs/python-modules/torch_tpu/default.nix
+++ b/nix-builder/pkgs/python-modules/torch_tpu/default.nix
@@ -2,7 +2,7 @@
   lib,
   autoPatchelfHook,
   buildPythonPackage,
-  fetchurl,
+  requireFile,
   python,
   stdenv,
 
@@ -23,27 +23,24 @@
   torch,
 }:
 
-# 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:
+# torch_tpu is served from Google's gated Artifact Registry, which
+# requires an OAuth2 bearer token. That token can't be handed to a build
+# sandbox without making the derivation impure, so instead the wheel is
+# prefetched out-of-band and handed to `requireFile`, which only checks
+# that a file matching `name` + `hash` already exists in the store:
 #
-#   export GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token)
+#   GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+#     nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu 0.1.1.dev20260707090224
 #
-# (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` below, then recompute `hash` with:
-#
-#   nix store prefetch-file --json --hash-type sha256 \
-#     --option netrc-file <(printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' "$(gcloud auth print-access-token)") \
-#     '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' \
-#     | jq -r .hash
+# The script prints the `hash` line to paste in below. The dev build is
+# dated; bump the date suffix in `version` below, then re-run the
+# script with the new version.
 
 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); if either moves, recompute
-  # via the command above.
+  # via the prefetch script above.
   abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
 in
 buildPythonPackage rec {
@@ -51,18 +48,15 @@ buildPythonPackage rec {
   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}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
+  src = requireFile {
+    name = "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
-        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
+    message = ''
+      torch_tpu is served from a gated Google Artifact Registry and
+      cannot be fetched by a pure Nix build. Fetch and register it with:
+
+        GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+          nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu ${version} ${abi}
     '';
   };
 
diff --git a/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh b/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh
new file mode 100755
index 000000000..31c9f9edb
--- /dev/null
+++ b/nix-builder/scripts/helpers/prefetch-tpu-wheel.sh
@@ -0,0 +1,40 @@
+#!/usr/bin/env bash
+# Fetch a wheel from Google's gated torch-tpu-virtual-registry and
+# register it in the Nix store under the exact (name, sha256) that
+# `requireFile` expects in pkgs/python-modules/{libtpu,torch_tpu}/default.nix.
+#
+# The fetch (and the OAuth2 token it needs) happens here, out-of-band,
+# so the actual Nix build stays pure and never touches the network.
+#
+# Usage:
+#   GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+#     scripts/helpers/prefetch-tpu-wheel.sh libtpu 0.0.43 [cp313]
+#   GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+#     scripts/helpers/prefetch-tpu-wheel.sh torch_tpu 0.1.1.dev20260707090224 [cp313]
+#
+# Prints a `hash = "sha256-...";` line to paste into the matching
+# default.nix when bumping a version. 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 nix (with the `nix-command`
+# experimental feature) and jq.
+set -euo pipefail
+
+pkg="${1:?usage: $0 <libtpu|torch_tpu> <version> [abi]}"
+version="${2:?usage: $0 <libtpu|torch_tpu> <version> [abi]}"
+abi="${3:-cp313}"
+: "${GCLOUD_ACCESS_TOKEN:?set GCLOUD_ACCESS_TOKEN to \$(gcloud auth print-access-token)}"
+
+registry_pkg="${pkg//_/-}"
+name="${pkg}-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl"
+url="https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/${registry_pkg}/${name}"
+netrc="$(mktemp)"
+trap 'rm -f "${netrc}"' EXIT
+printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' \
+  "${GCLOUD_ACCESS_TOKEN}" >"${netrc}"
+
+hash="$(nix store prefetch-file --json --hash-type sha256 --name "${name}" \
+  --option netrc-file "${netrc}" "${url}" | jq -r .hash)"
+
+echo "hash = \"${hash}\"; # ${abi}"

From 3de0b0d54eb207d55bb35327350148327870484c Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 14:53:19 +0000
Subject: [PATCH 20/32] nix-builder: regenerate build-variants.json without tpu

---
 nix-builder/build-variants.json | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/nix-builder/build-variants.json b/nix-builder/build-variants.json
index da668552c..7fa073eba 100644
--- a/nix-builder/build-variants.json
+++ b/nix-builder/build-variants.json
@@ -54,9 +54,6 @@
       "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",

From 4902e1485db8d31f0e95b36745723da0c6cfe77e Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Mon, 27 Jul 2026 15:01:12 +0000
Subject: [PATCH 21/32] docs: regenerate build-variants.md (tpu excluded from
 bundle cache)

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 docs/source/builder/build-variants.md | 5 -----
 1 file changed, 5 deletions(-)

diff --git a/docs/source/builder/build-variants.md b/docs/source/builder/build-variants.md
index 5c7011d76..bdc9033fb 100644
--- a/docs/source/builder/build-variants.md
+++ b/docs/source/builder/build-variants.md
@@ -62,10 +62,6 @@ 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`
@@ -81,5 +77,4 @@ one or more of the following variants:
 - `torch-cuda`
 - `torch-metal`
 - `torch-rocm`
-- `torch-tpu`
 - `torch-xpu`

From 8734a427a08c5cbed46c22c7c6775a270d46d04f Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 28 Jul 2026 10:59:08 +0000
Subject: [PATCH 22/32] tpu: use dash in torch_tpu python-depends declaration

Dependency declaration keys use dashes (tvm-ffi, nvidia-cutlass-dsl),
with the mapping translating to the underscored pkg/import name.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 examples/kernels/relu-tpu/build.toml      | 2 +-
 kernels-data/src/config/mod.rs            | 4 ++--
 kernels-data/src/python_dependencies.json | 2 +-
 kernels/src/kernels/python_depends.json   | 2 +-
 4 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/examples/kernels/relu-tpu/build.toml b/examples/kernels/relu-tpu/build.toml
index c930a6a61..5e2bd6909 100644
--- a/examples/kernels/relu-tpu/build.toml
+++ b/examples/kernels/relu-tpu/build.toml
@@ -9,7 +9,7 @@ backends = ["tpu"]
 repo-id = "kernels-test/relu-tpu"
 
 [general.tpu]
-python-depends = ["torch_tpu", "jax"]
+python-depends = ["torch-tpu", "jax"]
 
 [torch-noarch]
 
diff --git a/kernels-data/src/config/mod.rs b/kernels-data/src/config/mod.rs
index ed111d70d..1a6b98b15 100644
--- a/kernels-data/src/config/mod.rs
+++ b/kernels-data/src/config/mod.rs
@@ -476,7 +476,7 @@ mod tests {
             cuda: None,
             neuron: None,
             tpu: Some(TpuGeneral {
-                python_depends: Some(vec!["torch_tpu".to_string()]),
+                python_depends: Some(vec!["torch-tpu".to_string()]),
             }),
             xpu: None,
         }
@@ -491,7 +491,7 @@ mod tests {
             .collect::<Result<Vec<_>>>()
             .unwrap();
 
-        assert_eq!(deps, vec!["torch_tpu".to_string()]);
+        assert_eq!(deps, vec!["torch-tpu".to_string()]);
     }
 
     #[test]
diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json
index 8de880f5b..127c3faee 100644
--- a/kernels-data/src/python_dependencies.json
+++ b/kernels-data/src/python_dependencies.json
@@ -36,7 +36,7 @@
       }
     },
     "tpu": {
-      "torch_tpu": {
+      "torch-tpu": {
         "nix": ["torch_tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
       },
diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json
index 8de880f5b..127c3faee 100644
--- a/kernels/src/kernels/python_depends.json
+++ b/kernels/src/kernels/python_depends.json
@@ -36,7 +36,7 @@
       }
     },
     "tpu": {
-      "torch_tpu": {
+      "torch-tpu": {
         "nix": ["torch_tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
       },

From a4049f3ce9af5e1c1353959b49b7a140fcf5f5bc Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 28 Jul 2026 11:17:13 +0000
Subject: [PATCH 23/32] kernels-data: only add tpu general option to v5

We don't add new options to older build.toml editions once a newer
one is out, so v3/v4 no longer gain `general.tpu`; v5 is enough.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 kernels-data/src/config/v3.rs | 18 +-----------------
 kernels-data/src/config/v4.rs | 18 +-----------------
 2 files changed, 2 insertions(+), 34 deletions(-)

diff --git a/kernels-data/src/config/v3.rs b/kernels-data/src/config/v3.rs
index 766a0cf9f..b9aa400a6 100644
--- a/kernels-data/src/config/v3.rs
+++ b/kernels-data/src/config/v3.rs
@@ -50,8 +50,6 @@ pub struct General {
 
     pub python_depends: Option<Vec<String>>,
 
-    pub tpu: Option<TpuGeneral>,
-
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -69,12 +67,6 @@ 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 {
@@ -214,7 +206,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),
+            tpu: None,
             xpu: general.xpu.map(Into::into),
         })
     }
@@ -238,14 +230,6 @@ 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 {
diff --git a/kernels-data/src/config/v4.rs b/kernels-data/src/config/v4.rs
index ecc6c0125..1624cf983 100644
--- a/kernels-data/src/config/v4.rs
+++ b/kernels-data/src/config/v4.rs
@@ -49,8 +49,6 @@ pub struct General {
 
     pub python_depends: Option<Vec<String>>,
 
-    pub tpu: Option<TpuGeneral>,
-
     pub xpu: Option<XpuGeneral>,
 }
 
@@ -68,12 +66,6 @@ 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 {
@@ -212,7 +204,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),
+            tpu: None,
             xpu: general.xpu.map(Into::into),
         }
     }
@@ -248,14 +240,6 @@ 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 {

From fe2c74253b6379e1b9f66880f67845040a2f010e Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 28 Jul 2026 13:39:47 +0000
Subject: [PATCH 24/32] chore(nix-builder): remove empty dependencies from
 libtpu

---
 nix-builder/pkgs/python-modules/libtpu/default.nix | 2 --
 1 file changed, 2 deletions(-)

diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix
index ff18bdbab..bcad0e95e 100644
--- a/nix-builder/pkgs/python-modules/libtpu/default.nix
+++ b/nix-builder/pkgs/python-modules/libtpu/default.nix
@@ -52,8 +52,6 @@ buildPythonPackage rec {
     '';
   };
 
-  dependencies = [ ];
-
   pythonImportsCheck = [ "libtpu" ];
   doInstallCheck = false; # requires actual /dev/accel
 

From 9c8140a815167d3cd8ef55e63f1c1f68f5bbf2d3 Mon Sep 17 00:00:00 2001
From: Alvaro Moran <alvaro.moran@huggingface.co>
Date: Tue, 28 Jul 2026 14:43:50 +0000
Subject: [PATCH 25/32] kernels-data: drop jax/libtpu from tpu python-depends

Both are already hard, propagated Nix dependencies of torch_tpu
itself, so declaring them separately in python_dependencies.json
is redundant and get-kernel-check's sandbox still resolves them
transitively through torch-tpu.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
---
 examples/kernels/relu-tpu/build.toml      | 2 +-
 kernels-data/src/python_dependencies.json | 8 --------
 kernels/src/kernels/python_depends.json   | 8 --------
 3 files changed, 1 insertion(+), 17 deletions(-)

diff --git a/examples/kernels/relu-tpu/build.toml b/examples/kernels/relu-tpu/build.toml
index 5e2bd6909..0e5357235 100644
--- a/examples/kernels/relu-tpu/build.toml
+++ b/examples/kernels/relu-tpu/build.toml
@@ -9,7 +9,7 @@ backends = ["tpu"]
 repo-id = "kernels-test/relu-tpu"
 
 [general.tpu]
-python-depends = ["torch-tpu", "jax"]
+python-depends = ["torch-tpu"]
 
 [torch-noarch]
 
diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json
index 127c3faee..04339fd07 100644
--- a/kernels-data/src/python_dependencies.json
+++ b/kernels-data/src/python_dependencies.json
@@ -39,14 +39,6 @@
       "torch-tpu": {
         "nix": ["torch_tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
-      },
-      "jax": {
-        "nix": ["jax"],
-        "python": [{ "pkg": "jax", "import": "jax" }]
-      },
-      "libtpu": {
-        "nix": ["libtpu"],
-        "python": [{ "pkg": "libtpu" }]
       }
     }
   }
diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json
index 127c3faee..04339fd07 100644
--- a/kernels/src/kernels/python_depends.json
+++ b/kernels/src/kernels/python_depends.json
@@ -39,14 +39,6 @@
       "torch-tpu": {
         "nix": ["torch_tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
-      },
-      "jax": {
-        "nix": ["jax"],
-        "python": [{ "pkg": "jax", "import": "jax" }]
-      },
-      "libtpu": {
-        "nix": ["libtpu"],
-        "python": [{ "pkg": "libtpu" }]
       }
     }
   }

From ab334544fb044c18f04fc5e6943cf5644329c500 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Thu, 30 Jul 2026 18:45:08 +0000
Subject: [PATCH 26/32] python3Packages.torch-tpu: make import check work in
 the sandbox

Also canonicalize name.
---
 kernels-data/src/python_dependencies.json     |  2 +-
 kernels/src/kernels/python_depends.json       |  2 +-
 nix-builder/overlay.nix                       |  2 +-
 .../{torch_tpu => torch-tpu}/default.nix      | 29 +++++++++++++++----
 4 files changed, 26 insertions(+), 9 deletions(-)
 rename nix-builder/pkgs/python-modules/{torch_tpu => torch-tpu}/default.nix (82%)

diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json
index 04339fd07..1c2e5b6cd 100644
--- a/kernels-data/src/python_dependencies.json
+++ b/kernels-data/src/python_dependencies.json
@@ -37,7 +37,7 @@
     },
     "tpu": {
       "torch-tpu": {
-        "nix": ["torch_tpu"],
+        "nix": ["torch-tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
       }
     }
diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json
index 04339fd07..1c2e5b6cd 100644
--- a/kernels/src/kernels/python_depends.json
+++ b/kernels/src/kernels/python_depends.json
@@ -37,7 +37,7 @@
     },
     "tpu": {
       "torch-tpu": {
-        "nix": ["torch_tpu"],
+        "nix": ["torch-tpu"],
         "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
       }
     }
diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix
index 49d335ad3..13059f024 100644
--- a/nix-builder/overlay.nix
+++ b/nix-builder/overlay.nix
@@ -163,7 +163,7 @@ final: prev:
 
         libtpu = python-self.callPackage ./pkgs/python-modules/libtpu { };
 
-        torch_tpu = python-self.callPackage ./pkgs/python-modules/torch_tpu { };
+        torch-tpu = python-self.callPackage ./pkgs/python-modules/torch-tpu { };
 
         jupyter-server = python-super.jupyter-server.overrideAttrs (
           _: prevAttrs: {
diff --git a/nix-builder/pkgs/python-modules/torch_tpu/default.nix b/nix-builder/pkgs/python-modules/torch-tpu/default.nix
similarity index 82%
rename from nix-builder/pkgs/python-modules/torch_tpu/default.nix
rename to nix-builder/pkgs/python-modules/torch-tpu/default.nix
index c5dc007f9..182b3fe01 100644
--- a/nix-builder/pkgs/python-modules/torch_tpu/default.nix
+++ b/nix-builder/pkgs/python-modules/torch-tpu/default.nix
@@ -1,10 +1,12 @@
 {
   lib,
-  autoPatchelfHook,
+  stdenv,
   buildPythonPackage,
   requireFile,
+
+  autoPatchelfHook,
+  proot,
   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
@@ -44,7 +46,7 @@ let
   abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
 in
 buildPythonPackage rec {
-  pname = "torch_tpu";
+  pname = "torch-tpu";
   version = "0.1.1.dev20260707090224";
   format = "wheel";
 
@@ -73,14 +75,29 @@ buildPythonPackage rec {
   # 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 ];
+  nativeBuildInputs = [
+    autoPatchelfHook
+    proot
+  ];
   buildInputs = [ stdenv.cc.cc.lib ];
+
   preFixup = ''
     addAutoPatchelfSearchPath ${torch}/${python.sitePackages}/torch/lib
   '';
 
-  pythonImportsCheck = [ "torch_tpu" ];
-  doInstallCheck = false; # requires actual /dev/accel
+  installCheckPhase = ''
+    # tcmalloc hard-crashes if /sys is not available:
+    # 
+    # https://github.com/google/tcmalloc/issues/245
+    #
+    # We would still like to do an import check with sandboxing enabled,
+    # so use proot to fake presence of the relevant part of /sys.
+    mkdir -p fake-sys/devices/system/cpu
+    echo "0-1" > fake-sys/devices/system/cpu/possible
+    PYTHONPATH="$out/${python.sitePackages}:$PYTHONPATH" \
+      proot -b fake-sys:/sys \
+        ${python.interpreter} -c 'import torch_tpu'
+  '';
 
   meta = with lib; {
     description = "Torch TPU backend (PrivateUse1 name: \"tpu\")";

From 286a048b7a199d177b7bb76ec070552c19bd917b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Jul 2026 09:08:53 +0000
Subject: [PATCH 27/32] python3Packages.torch: fold torch-tpu into torch

---
 nix-builder/lib/mk-build-set.nix              |  1 +
 nix-builder/lib/variants/torch.nix            |  2 +
 .../pkgs/python-modules/hooks/default.nix     | 14 ++++++
 .../hooks/python-kernel-imports-check-hook.sh | 45 ++++++++++++++++++
 .../python-modules/torch/binary/generic.nix   | 47 ++++++++++++++++++-
 5 files changed, 108 insertions(+), 1 deletion(-)
 create mode 100644 nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh

diff --git a/nix-builder/lib/mk-build-set.nix b/nix-builder/lib/mk-build-set.nix
index 2b0355c63..74a4ff573 100644
--- a/nix-builder/lib/mk-build-set.nix
+++ b/nix-builder/lib/mk-build-set.nix
@@ -88,6 +88,7 @@ let
       # (unfree), so the tpu buildSet needs allowUnfree just like the
       # cuda/rocm/xpu sets. See pkgs/python-modules/libtpu/default.nix.
       allowUnfree = true;
+      tpuSupport = true;
     };
 
     xpu = {
diff --git a/nix-builder/lib/variants/torch.nix b/nix-builder/lib/variants/torch.nix
index 406a75fd4..7b9469277 100644
--- a/nix-builder/lib/variants/torch.nix
+++ b/nix-builder/lib/variants/torch.nix
@@ -13,6 +13,8 @@ let
       "metal"
     else if buildConfig.backend == "rocm" then
       "rocm${flattenVersion (lib.versions.majorMinor buildConfig.rocmVersion)}"
+    else if buildConfig.backend == "tpu" then
+      "tpu"
     else if buildConfig.backend == "xpu" then
       "xpu${flattenVersion (lib.versions.majorMinor buildConfig.xpuVersion)}"
     else
diff --git a/nix-builder/pkgs/python-modules/hooks/default.nix b/nix-builder/pkgs/python-modules/hooks/default.nix
index 246fc3682..32163a0b1 100644
--- a/nix-builder/pkgs/python-modules/hooks/default.nix
+++ b/nix-builder/pkgs/python-modules/hooks/default.nix
@@ -4,8 +4,22 @@ let
   inherit (python) pythonOnBuildForHost;
   pythonInterpreter = pythonOnBuildForHost.interpreter;
   pythonSitePackages = python.sitePackages;
+  pythonCheckInterpreter = python.interpreter;
 in
 {
+  # Ideally we'd just call this pythonImportsCheckHook, but it  would cause
+  # a rebuild of most nixpkgs dependencies.
+  pythonKernelImportsCheckHook = callPackage (
+    { makePythonHook }:
+    makePythonHook {
+      name = "python-kernel-imports-check-hook.sh";
+      substitutions = {
+        inherit pythonCheckInterpreter pythonSitePackages;
+        inherit (pkgs) proot;
+      };
+    } ./python-kernel-imports-check-hook.sh
+  ) { };
+
   pythonRelaxWheelDepsHook = callPackage (
     { makePythonHook, wheel }:
     makePythonHook {
diff --git a/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
new file mode 100644
index 000000000..f8150f69b
--- /dev/null
+++ b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
@@ -0,0 +1,45 @@
+# shellcheck shell=bash
+
+# This is the imports check hook from nixpkgs, modified to fake a /sys path
+# required for tcmalloc. Without this path, tcmalloc will crash:
+#
+#
+# https://github.com/google/tcmalloc/issues/245
+
+# Setup hook for checking whether Python imports succeed
+echo "Sourcing python-kernel-imports-check-hook.sh"
+
+pythonKernelImportsCheckPhase() {
+    echo "Executing pythonKernelImportsCheckPhase"
+
+    if [[ -n "${pythonKernelImportsCheck[*]-}" ]]; then
+        echo "Check whether the following modules can be imported: ${pythonKernelImportsCheck[*]}"
+        # shellcheck disable=SC2154
+        pythonKernelImportsCheckOutput="$out"
+        if [[ -n "${python-}" ]]; then
+            echo "Using python specific output \$python for imports check"
+            pythonKernelImportsCheckOutput=$python
+        fi
+        export PYTHONPATH="$pythonKernelImportsCheckOutput/@pythonSitePackages@:$PYTHONPATH"
+
+        # Prepare fake /sys
+        FAKESYS="$(mktemp -d)"
+        trap 'rm -rf -- "${FAKESYS}"' EXIT
+        mkdir -p "${FAKESYS}/devices/system/cpu"
+        echo "0-1" > "${FAKESYS}/devices/system/cpu/possible"
+
+        # Python modules and namespaces names are Python identifiers, which must not contain spaces.
+        # See https://docs.python.org/3/reference/lexical_analysis.html
+        # shellcheck disable=SC2048,SC2086
+        (
+          cd "$pythonKernelImportsCheckOutput" && 
+          @proot@/bin/proot -b ${FAKESYS}:/sys \
+          @pythonCheckInterpreter@ -c 'import sys; import importlib; list(map(lambda mod: importlib.import_module(mod), sys.argv[1:]))' ${pythonKernelImportsCheck[*]}
+        )
+    fi
+}
+
+if [[ -z "${dontUsePythonImportsCheck-}" ]]; then
+    echo "Using pythonKernelImportsCheckPhase"
+    appendToVar preDistPhases pythonKernelImportsCheckPhase
+fi
diff --git a/nix-builder/pkgs/python-modules/torch/binary/generic.nix b/nix-builder/pkgs/python-modules/torch/binary/generic.nix
index c66c6cb97..1b8183511 100644
--- a/nix-builder/pkgs/python-modules/torch/binary/generic.nix
+++ b/nix-builder/pkgs/python-modules/torch/binary/generic.nix
@@ -7,10 +7,12 @@
   buildPythonPackage,
   fetchFromGitHub,
   fetchurl,
+  requireFile,
 
   cudaSupport ? config.cudaSupport,
   metalSupport ? config.metalSupport,
   rocmSupport ? config.rocmSupport,
+  tpuSupport ? (config.tpuSupport or false),
   tritonSupport ? (!stdenv.hostPlatform.isDarwin),
   xpuSupport ? (config.xpuSupport or false),
 
@@ -18,6 +20,7 @@
   autoAddDriverRunpath,
   autoPatchelfHook,
   python,
+  pythonKernelImportsCheckHook,
   pythonRelaxWheelDepsHook,
   pythonWheelDepsCheckHook,
 
@@ -31,15 +34,21 @@
   # Python dependencies
   cuda-bindings,
   filelock,
+  frozendict,
   fsspec,
+  immutabledict,
+  jax,
   jinja2,
+  libtpu,
   networkx,
   numpy,
+  portpicker,
   pyyaml,
   pyzes,
   requests,
   setuptools,
   sympy,
+  tensorboard,
   triton,
   triton-cuda,
   triton-rocm,
@@ -122,6 +131,24 @@ let
     '';
   };
 
+  # Once we require multiple versions, factor this out.
+  torch-tpu =
+    let
+      abi = "cp${lib.versions.major python.version}${lib.versions.minor python.version}";
+      version = "0.1.1.dev20260707090224";
+    in
+    requireFile {
+      name = "torch_tpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
+      hash = "sha256-eCjwoX0UKd/L/IccxXn88p0GdyQjFdPl2Er1luZz4H0="; # cp313
+      message = ''
+        torch_tpu is served from a gated Google Artifact Registry and
+        cannot be fetched by a pure Nix build. Fetch and register it with:
+
+          GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
+            nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu ${version} ${abi}
+      '';
+    };
+
 in
 buildPythonPackage.override { stdenv = effectiveStdenv; } {
   pname = "torch";
@@ -139,10 +166,20 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
     inherit url hash;
   };
 
+  # If we are building for TPU, merge torch-tpu into the output path.
+  # Note that we cannot simply depend on torch-tpu, because torch-tpu
+  # itself depends on torch. The wheel unpacking hook does not seemt
+  # to support `srcs`, so we'll just copy over the wheel on TPU.
+  preUnpack = lib.optionals tpuSupport ''
+    mkdir -p dist
+    cp "${torch-tpu}" "dist/$(stripHash "${torch-tpu}")"
+  '';
+
   # We use our own pythonWheelDepsHook.
   dontCheckRuntimeDeps = true;
 
   nativeBuildInputs = [
+    pythonKernelImportsCheckHook
     pythonRelaxWheelDepsHook
     pythonWheelDepsCheckHook
   ]
@@ -225,6 +262,14 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
     sympy
     typing-extensions
   ]
+  ++ lib.optionals tpuSupport [
+    libtpu
+    jax
+    frozendict
+    immutabledict
+    portpicker
+    tensorboard
+  ]
   ++ lib.optionals tritonSupport [
     effectiveTriton
   ]
@@ -367,7 +412,7 @@ buildPythonPackage.override { stdenv = effectiveStdenv; } {
 
   dontStrip = true;
 
-  pythonImportsCheck = [ "torch" ];
+  pythonKernelImportsCheck = if tpuSupport then [ "torch_tpu" ] else [ "torch" ];
 
   passthru = {
     inherit

From 59dd5d41bb074c7b856f543d73bb05932f23e370 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Jul 2026 09:09:29 +0000
Subject: [PATCH 28/32] python3Packages.libtpu: 0.0.43 -> 0.0.44.1

---
 nix-builder/pkgs/python-modules/libtpu/default.nix | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nix-builder/pkgs/python-modules/libtpu/default.nix b/nix-builder/pkgs/python-modules/libtpu/default.nix
index bcad0e95e..e3cf3c198 100644
--- a/nix-builder/pkgs/python-modules/libtpu/default.nix
+++ b/nix-builder/pkgs/python-modules/libtpu/default.nix
@@ -37,12 +37,12 @@ let
 in
 buildPythonPackage rec {
   pname = "libtpu";
-  version = "0.0.43";
+  version = "0.0.44.1";
   format = "wheel";
 
   src = requireFile {
     name = "libtpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
-    hash = "sha256-X5LVmwJuRcNkMG+m7kKgcaOSKnWjeuw4HtCU+lfWamY="; # cp313
+    hash = "sha256-qd34JzD0TpphtqLj67/HqYgVg/bXT3gF6E+pYQyaCsA="; # cp313
     message = ''
       libtpu is served from a gated Google Artifact Registry and cannot
       be fetched by a pure Nix build. Fetch and register it with:

From 6e70381a5fbed2bb701f4432101ce1eab160aaa6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Jul 2026 09:12:08 +0000
Subject: [PATCH 29/32] Makefile: fixup `all` target

---
 Makefile | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Makefile b/Makefile
index 9f5239d88..b19055d87 100644
--- a/Makefile
+++ b/Makefile
@@ -3,7 +3,7 @@
 
 export check_dirs := kernels/src kernels/tests kernels-data/bindings/python
 
-all: src/kernels/python_depends.json
+all: kernels/src/kernels/python_depends.json
 
 kernels/src/kernels/python_depends.json: kernels-data/src/python_dependencies.json
 	cp $< $@

From 68676ad424b2f6b66ea1dd33e11491c0ab0ae1b0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Jul 2026 10:04:52 +0000
Subject: [PATCH 30/32] Remove torch-tpu and torch-tpu kernel dependency

---
 examples/kernels/relu-tpu/build.toml          |   3 -
 kernels-data/src/config/mod.rs                |  22 ++--
 kernels-data/src/python_dependencies.json     |   6 -
 kernels/src/kernels/python_depends.json       |   6 -
 nix-builder/overlay.nix                       |   2 -
 .../pkgs/python-modules/torch-tpu/default.nix | 108 ------------------
 6 files changed, 12 insertions(+), 135 deletions(-)
 delete mode 100644 nix-builder/pkgs/python-modules/torch-tpu/default.nix

diff --git a/examples/kernels/relu-tpu/build.toml b/examples/kernels/relu-tpu/build.toml
index 0e5357235..a51fee218 100644
--- a/examples/kernels/relu-tpu/build.toml
+++ b/examples/kernels/relu-tpu/build.toml
@@ -8,9 +8,6 @@ backends = ["tpu"]
 [general.hub]
 repo-id = "kernels-test/relu-tpu"
 
-[general.tpu]
-python-depends = ["torch-tpu"]
-
 [torch-noarch]
 
 [kernel]
diff --git a/kernels-data/src/config/mod.rs b/kernels-data/src/config/mod.rs
index 1a6b98b15..93ba480b0 100644
--- a/kernels-data/src/config/mod.rs
+++ b/kernels-data/src/config/mod.rs
@@ -463,7 +463,7 @@ pub enum ConfigError {
 mod tests {
     use super::*;
 
-    fn general_with_tpu_deps() -> General {
+    fn general_with_cuda_deps() -> General {
         General {
             name: KernelName::new("test-kernel").unwrap(),
             version: 1,
@@ -473,30 +473,32 @@ mod tests {
             backends: vec![Backend::Tpu],
             hub: None,
             python_depends: None,
-            cuda: None,
-            neuron: None,
-            tpu: Some(TpuGeneral {
-                python_depends: Some(vec!["torch-tpu".to_string()]),
+            cuda: Some(CudaGeneral {
+                minver: None,
+                maxver: None,
+                python_depends: Some(vec!["nvidia-cutlass-dsl".to_string()]),
             }),
+            neuron: None,
+            tpu: None,
             xpu: None,
         }
     }
 
     #[test]
-    fn backend_python_depends_resolves_tpu() {
-        let general = general_with_tpu_deps();
+    fn backend_python_depends_resolves_cuda() {
+        let general = general_with_cuda_deps();
         let deps = general
-            .backend_python_depends(Backend::Tpu)
+            .backend_python_depends(Backend::Cuda)
             .map(|dep| dep.map(|(name, _)| name.to_string()))
             .collect::<Result<Vec<_>>>()
             .unwrap();
 
-        assert_eq!(deps, vec!["torch-tpu".to_string()]);
+        assert_eq!(deps, vec!["nvidia-cutlass-dsl".to_string()]);
     }
 
     #[test]
     fn backend_python_depends_empty_for_backend_without_deps() {
-        let general = general_with_tpu_deps();
+        let general = general_with_cuda_deps();
 
         assert!(
             general
diff --git a/kernels-data/src/python_dependencies.json b/kernels-data/src/python_dependencies.json
index 1c2e5b6cd..266904e82 100644
--- a/kernels-data/src/python_dependencies.json
+++ b/kernels-data/src/python_dependencies.json
@@ -34,12 +34,6 @@
         "nix": [],
         "python": [{ "pkg": "onednn-devel" }]
       }
-    },
-    "tpu": {
-      "torch-tpu": {
-        "nix": ["torch-tpu"],
-        "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
-      }
     }
   }
 }
diff --git a/kernels/src/kernels/python_depends.json b/kernels/src/kernels/python_depends.json
index 1c2e5b6cd..266904e82 100644
--- a/kernels/src/kernels/python_depends.json
+++ b/kernels/src/kernels/python_depends.json
@@ -34,12 +34,6 @@
         "nix": [],
         "python": [{ "pkg": "onednn-devel" }]
       }
-    },
-    "tpu": {
-      "torch-tpu": {
-        "nix": ["torch-tpu"],
-        "python": [{ "pkg": "torch_tpu", "import": "torch_tpu" }]
-      }
     }
   }
 }
diff --git a/nix-builder/overlay.nix b/nix-builder/overlay.nix
index 13059f024..bb49d2fa0 100644
--- a/nix-builder/overlay.nix
+++ b/nix-builder/overlay.nix
@@ -163,8 +163,6 @@ final: prev:
 
         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/torch-tpu/default.nix b/nix-builder/pkgs/python-modules/torch-tpu/default.nix
deleted file mode 100644
index 182b3fe01..000000000
--- a/nix-builder/pkgs/python-modules/torch-tpu/default.nix
+++ /dev/null
@@ -1,108 +0,0 @@
-{
-  lib,
-  stdenv,
-  buildPythonPackage,
-  requireFile,
-
-  autoPatchelfHook,
-  proot,
-  python,
-
-  # 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,
-}:
-
-# torch_tpu is served from Google's gated Artifact Registry, which
-# requires an OAuth2 bearer token. That token can't be handed to a build
-# sandbox without making the derivation impure, so instead the wheel is
-# prefetched out-of-band and handed to `requireFile`, which only checks
-# that a file matching `name` + `hash` already exists in the store:
-#
-#   GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
-#     nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu 0.1.1.dev20260707090224
-#
-# The script prints the `hash` line to paste in below. The dev build is
-# dated; bump the date suffix in `version` below, then re-run the
-# script with the new version.
-
-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); if either moves, recompute
-  # via the prefetch script above.
-  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 = requireFile {
-    name = "torch_tpu-${version}-${abi}-${abi}-manylinux_2_31_x86_64.whl";
-    hash = "sha256-eCjwoX0UKd/L/IccxXn88p0GdyQjFdPl2Er1luZz4H0="; # cp313
-    message = ''
-      torch_tpu is served from a gated Google Artifact Registry and
-      cannot be fetched by a pure Nix build. Fetch and register it with:
-
-        GCLOUD_ACCESS_TOKEN=$(gcloud auth print-access-token) \
-          nix-builder/scripts/helpers/prefetch-tpu-wheel.sh torch_tpu ${version} ${abi}
-    '';
-  };
-
-  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
-    proot
-  ];
-  buildInputs = [ stdenv.cc.cc.lib ];
-
-  preFixup = ''
-    addAutoPatchelfSearchPath ${torch}/${python.sitePackages}/torch/lib
-  '';
-
-  installCheckPhase = ''
-    # tcmalloc hard-crashes if /sys is not available:
-    # 
-    # https://github.com/google/tcmalloc/issues/245
-    #
-    # We would still like to do an import check with sandboxing enabled,
-    # so use proot to fake presence of the relevant part of /sys.
-    mkdir -p fake-sys/devices/system/cpu
-    echo "0-1" > fake-sys/devices/system/cpu/possible
-    PYTHONPATH="$out/${python.sitePackages}:$PYTHONPATH" \
-      proot -b fake-sys:/sys \
-        ${python.interpreter} -c 'import torch_tpu'
-  '';
-
-  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" ];
-  };
-}

From 9864d91058e092719bf79391fd0374211973f243 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Jul 2026 10:05:29 +0000
Subject: [PATCH 31/32] Make `get-kernel-check` hook work on TPU

---
 nix-builder/pkgs/get-kernel-check/default.nix        |  7 ++++++-
 .../pkgs/get-kernel-check/get-kernel-check-hook.sh   | 12 ++++++++++++
 nix-builder/pkgs/python-modules/hooks/default.nix    |  2 +-
 .../hooks/python-kernel-imports-check-hook.sh        |  5 ++---
 4 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/nix-builder/pkgs/get-kernel-check/default.nix b/nix-builder/pkgs/get-kernel-check/default.nix
index 51e492832..a9ac4cb67 100644
--- a/nix-builder/pkgs/get-kernel-check/default.nix
+++ b/nix-builder/pkgs/get-kernel-check/default.nix
@@ -1,8 +1,13 @@
-{ makeSetupHook, python3 }:
+{
+  makeSetupHook,
+  proot,
+  python3,
+}:
 
 makeSetupHook {
   name = "get-kernel-check-hook";
   substitutions = {
+    proot = "${proot}/bin/proot";
     python3 = "${python3}/bin/python";
     kernels = "${with python3.pkgs; makePythonPath [ kernels ]}";
   };
diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh
index 23d4501a7..61c7a65b0 100755
--- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh
+++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh
@@ -21,11 +21,23 @@ _getKernelCheckHook() {
     export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}"
   fi
 
+
+  # Prepare fake /sys for tcmalloc. Without this path, tcmalloc will crash:
+  #
+  # https://github.com/google/tcmalloc/issues/245
+  #
+  # tcmalloc is used by the TPU libraries.
+  FAKESYS="$(mktemp -d)"
+  trap 'rm -rf -- "${FAKESYS}"' EXIT
+  mkdir -p "${FAKESYS}/devices/system/cpu"
+  echo "0-1" > "${FAKESYS}/devices/system/cpu/possible"
+
   # Some kernels want to write stuff (especially when they use Triton).
   HOME=$(mktemp -d -t test.XXXXXX) || exit 1
   trap "rm -rf '$HOME'" EXIT
 
   PYTHONPATH="@kernels@" \
+    @proot@ -b ${FAKESYS}:/sys \
     @python3@ -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${out}'))"
 }
 
diff --git a/nix-builder/pkgs/python-modules/hooks/default.nix b/nix-builder/pkgs/python-modules/hooks/default.nix
index 32163a0b1..7343dde47 100644
--- a/nix-builder/pkgs/python-modules/hooks/default.nix
+++ b/nix-builder/pkgs/python-modules/hooks/default.nix
@@ -15,7 +15,7 @@ in
       name = "python-kernel-imports-check-hook.sh";
       substitutions = {
         inherit pythonCheckInterpreter pythonSitePackages;
-        inherit (pkgs) proot;
+        proot = "${pkgs.proot}/bin/proot";
       };
     } ./python-kernel-imports-check-hook.sh
   ) { };
diff --git a/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
index f8150f69b..ceee0a4fb 100644
--- a/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
+++ b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
@@ -3,7 +3,6 @@
 # This is the imports check hook from nixpkgs, modified to fake a /sys path
 # required for tcmalloc. Without this path, tcmalloc will crash:
 #
-#
 # https://github.com/google/tcmalloc/issues/245
 
 # Setup hook for checking whether Python imports succeed
@@ -22,7 +21,7 @@ pythonKernelImportsCheckPhase() {
         fi
         export PYTHONPATH="$pythonKernelImportsCheckOutput/@pythonSitePackages@:$PYTHONPATH"
 
-        # Prepare fake /sys
+        # Prepare fake /sys for tcmalloc.
         FAKESYS="$(mktemp -d)"
         trap 'rm -rf -- "${FAKESYS}"' EXIT
         mkdir -p "${FAKESYS}/devices/system/cpu"
@@ -33,7 +32,7 @@ pythonKernelImportsCheckPhase() {
         # shellcheck disable=SC2048,SC2086
         (
           cd "$pythonKernelImportsCheckOutput" && 
-          @proot@/bin/proot -b ${FAKESYS}:/sys \
+          @proot@ -b ${FAKESYS}:/sys \
           @pythonCheckInterpreter@ -c 'import sys; import importlib; list(map(lambda mod: importlib.import_module(mod), sys.argv[1:]))' ${pythonKernelImportsCheck[*]}
         )
     fi

From 40427a75621173bddbed5987b989a5e2f36b1cc8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Fri, 31 Jul 2026 12:02:09 +0000
Subject: [PATCH 32/32] nix-builder: make /sys faking in hooks conditional on
 `tpuSupport`

---
 nix-builder/pkgs/get-kernel-check/default.nix | 10 ++++++++-
 .../get-kernel-check/get-kernel-check-hook.sh | 22 +++++++++++--------
 .../pkgs/python-modules/hooks/default.nix     |  4 +++-
 .../hooks/python-kernel-imports-check-hook.sh | 19 ++++++++++------
 4 files changed, 37 insertions(+), 18 deletions(-)

diff --git a/nix-builder/pkgs/get-kernel-check/default.nix b/nix-builder/pkgs/get-kernel-check/default.nix
index a9ac4cb67..dc1d8eb98 100644
--- a/nix-builder/pkgs/get-kernel-check/default.nix
+++ b/nix-builder/pkgs/get-kernel-check/default.nix
@@ -1,14 +1,22 @@
 {
+  config,
+  lib,
   makeSetupHook,
+  stdenv,
+
   proot,
   python3,
 }:
 
+let
+  useFakeSys = config.tpuSupport or false;
+in
 makeSetupHook {
   name = "get-kernel-check-hook";
   substitutions = {
-    proot = "${proot}/bin/proot";
     python3 = "${python3}/bin/python";
     kernels = "${with python3.pkgs; makePythonPath [ kernels ]}";
+    proot = lib.optionalString useFakeSys "${proot}/bin/proot";
+    useFakeSys = lib.optionalString useFakeSys "1";
   };
 } ./get-kernel-check-hook.sh
diff --git a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh
index 61c7a65b0..5d229f8ef 100755
--- a/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh
+++ b/nix-builder/pkgs/get-kernel-check/get-kernel-check-hook.sh
@@ -21,23 +21,27 @@ _getKernelCheckHook() {
     export DYLD_LIBRARY_PATH="${TORCH_DIR}/lib:${DYLD_LIBRARY_PATH}"
   fi
 
+  HOME=$(mktemp -d -t test.XXXXXX) || exit 1
+  trap "rm -rf '$HOME'" EXIT
 
   # Prepare fake /sys for tcmalloc. Without this path, tcmalloc will crash:
   #
   # https://github.com/google/tcmalloc/issues/245
   #
   # tcmalloc is used by the TPU libraries.
-  FAKESYS="$(mktemp -d)"
-  trap 'rm -rf -- "${FAKESYS}"' EXIT
-  mkdir -p "${FAKESYS}/devices/system/cpu"
-  echo "0-1" > "${FAKESYS}/devices/system/cpu/possible"
-
-  # Some kernels want to write stuff (especially when they use Triton).
-  HOME=$(mktemp -d -t test.XXXXXX) || exit 1
-  trap "rm -rf '$HOME'" EXIT
+  local prootCmd=""
+  if [[ -n "@useFakeSys@" ]]; then
+      echo "Faking /sys for tcmalloc"
+      local fakeSys
+      fakeSys="$(mktemp -d)"
+      trap 'rm -rf -- "${fakeSys}"' EXIT
+      mkdir -p "${fakeSys}/devices/system/cpu"
+      echo "0-1" > "${fakeSys}/devices/system/cpu/possible"
+      prootCmd="@proot@ -b ${fakeSys}:/sys"
+  fi
 
   PYTHONPATH="@kernels@" \
-    @proot@ -b ${FAKESYS}:/sys \
+    ${prootCmd} \
     @python3@ -c "from pathlib import Path; import kernels; kernels.get_local_kernel(Path('${out}'))"
 }
 
diff --git a/nix-builder/pkgs/python-modules/hooks/default.nix b/nix-builder/pkgs/python-modules/hooks/default.nix
index 7343dde47..f4eae33ef 100644
--- a/nix-builder/pkgs/python-modules/hooks/default.nix
+++ b/nix-builder/pkgs/python-modules/hooks/default.nix
@@ -5,6 +5,7 @@ let
   pythonInterpreter = pythonOnBuildForHost.interpreter;
   pythonSitePackages = python.sitePackages;
   pythonCheckInterpreter = python.interpreter;
+  useFakeSys = pkgs.config.tpuSupport or false;
 in
 {
   # Ideally we'd just call this pythonImportsCheckHook, but it  would cause
@@ -15,7 +16,8 @@ in
       name = "python-kernel-imports-check-hook.sh";
       substitutions = {
         inherit pythonCheckInterpreter pythonSitePackages;
-        proot = "${pkgs.proot}/bin/proot";
+        proot = lib.optionalString useFakeSys "${pkgs.proot}/bin/proot";
+        useFakeSys = lib.optionalString useFakeSys "1";
       };
     } ./python-kernel-imports-check-hook.sh
   ) { };
diff --git a/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
index ceee0a4fb..49e41ae86 100644
--- a/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
+++ b/nix-builder/pkgs/python-modules/hooks/python-kernel-imports-check-hook.sh
@@ -21,18 +21,23 @@ pythonKernelImportsCheckPhase() {
         fi
         export PYTHONPATH="$pythonKernelImportsCheckOutput/@pythonSitePackages@:$PYTHONPATH"
 
-        # Prepare fake /sys for tcmalloc.
-        FAKESYS="$(mktemp -d)"
-        trap 'rm -rf -- "${FAKESYS}"' EXIT
-        mkdir -p "${FAKESYS}/devices/system/cpu"
-        echo "0-1" > "${FAKESYS}/devices/system/cpu/possible"
+        local prootCmd=""
+        if [[ -n "@useFakeSys@" ]]; then
+            echo "Faking /sys for tcmalloc"
+            local fakeSys
+            fakeSys="$(mktemp -d)"
+            trap 'rm -rf -- "${fakeSys}"' EXIT
+            mkdir -p "${fakeSys}/devices/system/cpu"
+            echo "0-1" > "${fakeSys}/devices/system/cpu/possible"
+            prootCmd="@proot@ -b ${fakeSys}:/sys"
+        fi
 
         # Python modules and namespaces names are Python identifiers, which must not contain spaces.
         # See https://docs.python.org/3/reference/lexical_analysis.html
         # shellcheck disable=SC2048,SC2086
         (
-          cd "$pythonKernelImportsCheckOutput" && 
-          @proot@ -b ${FAKESYS}:/sys \
+          cd "$pythonKernelImportsCheckOutput" &&
+          ${prootCmd} \
           @pythonCheckInterpreter@ -c 'import sys; import importlib; list(map(lambda mod: importlib.import_module(mod), sys.argv[1:]))' ${pythonKernelImportsCheck[*]}
         )
     fi