Add TPU backend support (torch_tpu / Pallas kernels) by tengomucho · Pull Request #714 · huggingface/kernels

tengomucho · 2026-07-13T21:56:37Z

Important: Read before submitting

Closes #711.

Adds TPU as a first-class backend across the stack, so a kernel repo declaring backends = ["tpu"] builds and loads a torch-tpu variant (Pallas/JAX kernels running through torch_tpu):

kernels: TPU backend/device detection via torch.backends.tpu, torch-tpu noarch variant resolution, kernelize device support, tpu_only pytest marker.
kernels-data: Backend::Tpu enum variant with serde/pyo3 bindings and TPU python_depends in backend metadata.
kernel-builder: tpu accepted in init --backends, TPU detection in the noarch _ops.py template.
nix-builder: TPU build set (jax, jaxlib, libtpu, torch_tpu), bundle/CI outputs, and the torch-tpu build variant. This part was made and tested with the help of Fable.
examples: relu-tpu example kernel repository.

TPU support is experimental (Tier 3): the libtpu/torch_tpu wheels are not public yet and are fetched from Google's Artifact Registry using a gcloud token (GCLOUD_ACCESS_TOKEN), so relu-tpu is excluded from the examples CI for now (until we have access to the TPUs for CI).

Note: KNOWN_BACKENDS in kernels.utils intentionally does not gain tpu, any value change hits an API compatibility lint error in the CI.

Checklist

This PR is linked to an issue that was discussed and approved
I have tested these changes locally
New/changed functionality has test coverage
LLM disclosure:
- I did not use an LLM to create this PR.
- I used and LLM for assistance while creating this PR.
- This PR was mostly or completely generated by an LLM.

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.

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.

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.

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.

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.

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).

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>

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>

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>

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>

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>

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>

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>

HuggingFaceDocBuilderDev · 2026-07-13T21:58:06Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

github-actions · 2026-07-13T23:18:19Z

Coverage report — `kernels/`

Measured on: Python 3.10 / Torch 2.12.0.
Other CI configurations are not included in this number.
Hardware-gated code paths (ROCm/XPU/NPU/Darwin/Windows) are excluded or unreachable on the Linux+CUDA runner.

Total coverage: 84.0% — threshold: 80% — ✅

Per-file breakdown

Name	Stmts	Miss	Cover	Missing
src/kernels/__init__.py	10	0	100%
src/kernels/_system.py	6	1	83%	10
src/kernels/_versions.py	63	7	89%	46, 49, 52-53, 56-57, 100
src/kernels/backends.py	212	62	71%	40, 44, 48-51, 68, 90, 108, 117, 121, 125-127, 148, 157, 161, 165-167, 188, 199, 201, 208-211, 224, 228, 232-252, 260, 283-303
src/kernels/compat.py	8	1	88%	5
src/kernels/deps.py	54	4	93%	58-59, 95, 98
src/kernels/layer/__init__.py	6	0	100%
src/kernels/layer/_interval_tree.py	103	4	96%	23, 52, 147, 150
src/kernels/layer/device.py	48	14	71%	42, 47-49, 91, 96-98, 101, 149, 152, 155-157
src/kernels/layer/func.py	81	7	91%	81, 111, 183, 301, 307, 320, 338
src/kernels/layer/globals.py	5	0	100%
src/kernels/layer/kernelize.py	74	8	89%	255, 281, 289-290, 296, 300, 316-318
src/kernels/layer/layer.py	210	16	92%	167, 210, 216, 229, 337, 417-418, 430, 439, 447, 458, 487, 491, 504, 557, 587
src/kernels/layer/mode.py	14	0	100%
src/kernels/layer/repos.py	144	42	71%	27, 33, 36-43, 63-64, 70, 73-76, 90, 94, 103-104, 110, 113-116, 123-124, 130, 133-136, 143-144, 150, 153-156, 163-164, 170, 173-176, 257
src/kernels/lockfile.py	71	46	35%	37-104, 108-131
src/kernels/status.py	49	2	96%	23, 81
src/kernels/utils.py	313	55	82%	65, 77-81, 87-88, 246, 250, 253, 315, 323, 362-363, 401, 432, 437, 472, 701, 704, 706, 712, 725-726, 747-759, 763-770, 778, 782-792, 796-803, 841, 845, 864, 866
src/kernels/variants.py	270	22	92%	63, 94, 115, 145, 254-255, 297-300, 302, 386-393, 399-405, 436-442, 454-460, 559-561
src/kernels/verify.py	88	1	99%	32
TOTAL	1829	292	84%

Updated by the Test kernels workflow on commit 4902e1485db8d31f0e95b36745723da0c6cfe77e.

drbh

looking good! just left some comments around fetching deps and the kernel-builder init support

drbh · 2026-07-15T23:08:39Z

    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`).


if we want to add support to initialize a tpu repo, I think we need to update build_init_fileset to handle has_tpu and render a tpu/noarch scaffold, including [general.tpu], [torch-noarch], and a tpu-specific Python module template. (similar to the other backends). otherwise we should probably leave it out of this list

Slightly related: neuron is already in that list and also has no has_neuron or template. Maybe we should remove it from here.

You are right @sayakpaul : but I think the neuron part should be addressed on another issue and PR.

drbh · 2026-07-15T23:08:53Z

+      printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword %s\n' \
+        "$GCLOUD_ACCESS_TOKEN" > netrc
+    '';
+  };


would it be possible to fetch this from pypi? I believe this is the library https://pypi.org/project/libtpu, maybe we can use fetchPypi?

The version available in pypi is an old version, Google hosts newerlibtpu version on an internal index reachable with token, and torch_tpu depends on those (at least for now).

drbh · 2026-07-15T23:08:57Z

@@ -0,0 +1,35 @@
+#!/usr/bin/env bash
+# Print the wheel URL and SRI hash for a torch-tpu-virtual-registry


it would probably be best to avoid adding this file, I think we only need it since the torch_tpu library requires auth.

in the case we move libtpu to pypi and only fetch torch_tpu from the google registry - I think we might be able to add the refresh command as a comment to the torch_tpu nix file? we might be able to get the hash with a one liner like below - and that would remove the need for this helper

nix store prefetch-file \ --json \ --hash-type sha256 \ --option netrc-file =({ printf 'machine us-python.pkg.dev\nlogin oauth2accesstoken\npassword '; gcloud auth print-access-token; }) \ 'https://us-python.pkg.dev/ml-oss-artifacts-transient/torch-tpu-virtual-registry/torch-tpu/torch_tpu-0.1.1.dev20260707090224-cp313-cp313-manylinux_2_31_x86_64.whl' \ | jq -r .hash

I just pushed a commit that addressed you comment.

drbh · 2026-07-15T23:08:59Z

+    tpu = true;
+    systems = [ "x86_64-linux" ];
+    bundleBuild = true;
+  }


would it be possible to also add entries 2.12 and 2.13 (not sure if there is a specific reason we only target 2.11)

Currently, torch_tpu is only compatible with 2.11, otherwise there is an ABI error (see https://github.com/google-pytorch/torch_tpu/issues/1533 if you have access)

Oh, oops, we will be removing 2.11 in #732 . Let me add enough of a 2.11 'skeleton' to make this work (basically setting other tritons, etc. to null).

sayakpaul

Nothing much from my side. I will have to defer the nix related bits to @danieldk @drbh. But the non-Nix related bits look good to me.

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.

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.

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.

Fixes CI's "Check docs are up to date" job, which went stale after tpu and neuron were dropped from the --backends help text.

danieldk · 2026-07-27T13:16:29Z

+    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
+        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


This requires an impure build though, we should avoid that people start doing impure builds, so I would remove this. Instead it's probably best to use requireFile:

https://nixos.org/manual/nixpkgs/stable/#requirefile

And then add a utility script that uses prefetch-file using GCLOUD_ACCESS_TOKEN.

In CI we also have to make sure that the output path doesn't get added to our cache, or we are publicly distributing it. I think, but am not 100%, sure that it does not get uploaded to the cache if the prefetch is done before installing Cachix (since IIRC we have use Cachix in store monitoring mode).

By the way, is there any chance that we can convince Google to make this available without a token?

Agreed on both. I'll switch to requireFile + a prefetch helper script for the actual fetch (similar to the one I have just removed after @drbh comment), so builds stay pure.
I also checked about the cache point, and tpu is currently marked bundleBuild = true, which puts it in .#forCache, built right after cachix-action installs its store-watcher in update_cache.yaml. Reordering the prefetch before Cachix would keep the raw wheel out, but torch_tpu's autoPatchelfHook step means the built package is itself a modified copy of gated content, so I don't think step-ordering alone is safe. I propose to leave bundleBuild = false for tpu for now so it's excluded from the shared cache pipeline entirely, until we get public packages for this.

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>

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>

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

danieldk

Great work! Added a bunch of comments.

danieldk · 2026-07-28T09:51:17Z

  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`, `rocm`, `xpu`)


TPU backend support is in (detection, Backend::Tpu, Nix build variant, relu-tpu example). What's not here is kernel-builder init scaffolding. TPU kernels are pure Python/JAX (torch-noarch + [general.tpu]), not C++ torch-ext bindings like init generates, so advertising --backends tpu there would scaffold the wrong thing. Dropped it from init's help for now, we can add it later when/if we define a real has_tpu template path.

Note the same is valid for neuron

danieldk · 2026-07-28T09:56:23Z

+repo-id = "kernels-test/relu-tpu"
+
+[general.tpu]
+python-depends = ["torch_tpu", "jax"]


Shouldn't we assume that at least torch_tpu is always installed when the tpu backend is used? Maybe the same for Jax, not 100% sure. In general, we should stop adding more Python dependencies.

I think this should largely be encapsulated in the loading logic of kernels itself, similar to how tvm-ffi variants are considered when tvm_ffi can be imported.

We use dashes in our dependency declarations (we can let the mapping take care of translating to underscores.

Agreed on jax and the naming. Fixed the dash issue (renamed the key to torch-tpu, still maps to the torch_tpu pkg/import).

On torch_tpu: the runtime check in kernels is indeed dead code (backend detection already requires torch_tpu to be present). But the same declaration also drives the nix-side dependency resolution in nix-builder/lib/deps.nix. That decides what's installed in the sandbox for get-kernel-check to import the kernel, nothing else pulls torch_tpu in there, so dropping it would break the nix build. Kept it for that reason.

But TPU has it's own build variant (tpu = true;), so we can conditionally include any dependencies needed in the Torch derivation. It's the same principle as all the other backends, except rather than having a backend-specific Torch build, we have the CPU build with additional dependencies.

Now the outcome is quite weird. If you would build .#torch.torch212-cxx11-tpu-x86_64-linux in the top-level kernels repo, you would actually get a CPU-only build, not an XPU-enabled build.

Agreed, torch_tpu's presence is already implied by backend detection (hasattr(torch.backends, "tpu") only goes true once torch_tpu is autoloaded), so declaring it as a dependency is redundant, and it's not on public PyPI anyway. Probably the same for libtpu. jax is the exception since it's public and not implied by detection. See this comment for why this isn't a trivial deletion though.

danieldk · 2026-07-28T09:57:48Z

    pub name: Option<RepoInfo>,

-    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `neuron`, `rocm`, `xpu`).
+    /// Backends to enable (`all`, `cpu`, `cuda`, `metal`, `rocm`, `xpu`).


Shouldn't tpu be added here?

Also, why is neuron removed? (Also see the removal of neuron at the start.)

I explained it in this comment.

danieldk · 2026-07-28T10:00:18Z


    pub python_depends: Option<Vec<String>>,

+    pub tpu: Option<TpuGeneral>,


We don't add options to older build.toml versions once a new one is out, so adding to v5 is enough.

danieldk · 2026-07-28T10:02:04Z

+    },
+    "tpu": {
+      "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" }]
+      }


I think of all these should be tested in kernels compat instead. Possibly with the exception of Jax.

We might also want to support pure Jax variants like jax-tpu, but that might be better in a follow-up PR.

Right in principle, but one wrinkle: this file is also read directly by nix-builder/lib/deps.nix to pick the nix packages for the build sandbox (used by get-kernel-check), not just to generate the pip pyproject.toml extra. So we can't just drop torch-tpu/libtpu here without breaking the nix build. Doing this properly means splitting "nix build dep" from "pip-installable dep" in the schema. If we want to split that, alongside with jax variants, I would do it in another Issue+PR, wdyt?

danieldk · 2026-07-28T10:05:42Z

+    '';
+  };
+
+  dependencies = [ ];


Not used, so can be removed.

danieldk · 2026-07-28T10:07:14Z

+  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


Aha, so if Jax always gets pulled in when importing torch_tpu, we also do not want Jax in python_dependencies.json.

OK, I dropped both jax and libtpu from python_dependencies.json's tpu block (only torch-tpu remains). Both are already hard, propagated Nix deps of torch_tpu itself. Verified end-to-end with a real build (nix build .../relu-tpu#packages.x86_64-linux.redistributable."torch-tpu", gated wheels fetched via prefetch-tpu-wheel.sh): torch_tpu still imports fine (jax/libtpu pulled in transitively), and get-kernel-check loads the built relu_tpu kernel successfully. No regression.

danieldk · 2026-07-28T10:09:50Z

+    tpu = true;
+    systems = [ "x86_64-linux" ];
+    bundleBuild = true;
+  }


Oh, oops, we will be removing 2.11 in #732 . Let me add enough of a 2.11 'skeleton' to make this work (basically setting other tritons, etc. to null).

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>

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>

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>

tengomucho and others added 13 commits July 13, 2026 15:37

examples: add flake.nix to relu-tpu kernel

38712a1

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/relu-tpu: annotate the jax function signature

7126d26

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>

tengomucho requested review from drbh and sayakpaul July 13, 2026 21:56

drbh requested changes Jul 15, 2026

View reviewed changes

drbh reviewed Jul 15, 2026

View reviewed changes

sayakpaul reviewed Jul 16, 2026

View reviewed changes

tengomucho added 2 commits July 27, 2026 11:45

Merge branch 'main' into tpu

38a0df3

tengomucho temporarily deployed to testpypi July 27, 2026 09:53 — with GitHub Actions Inactive

tengomucho added 2 commits July 27, 2026 12:05

tengomucho temporarily deployed to testpypi July 27, 2026 10:19 — with GitHub Actions Inactive

docs: regenerate builder-cli.md for init --backends help text

6f1e3a5

Fixes CI's "Check docs are up to date" job, which went stale after tpu and neuron were dropped from the --backends help text.

tengomucho temporarily deployed to testpypi July 27, 2026 10:49 — with GitHub Actions Inactive

tengomucho requested a review from drbh July 27, 2026 13:04

danieldk reviewed Jul 27, 2026

View reviewed changes

tengomucho and others added 2 commits July 27, 2026 14:18

tengomucho temporarily deployed to testpypi July 27, 2026 14:50 — with GitHub Actions Inactive

nix-builder: regenerate build-variants.json without tpu

3de0b0d

tengomucho temporarily deployed to testpypi July 27, 2026 14:54 — with GitHub Actions Inactive

docs: regenerate build-variants.md (tpu excluded from bundle cache)

4902e14

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tengomucho temporarily deployed to testpypi July 27, 2026 15:01 — with GitHub Actions Inactive

danieldk reviewed Jul 28, 2026

View reviewed changes

tengomucho and others added 2 commits July 28, 2026 10:59

tpu: use dash in torch_tpu python-depends declaration

8734a42

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>

kernels-data: only add tpu general option to v5

a4049f3

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>

danieldk mentioned this pull request Jul 28, 2026

Update to final Torch 2.13.0, cleanup #732

Open

tengomucho and others added 2 commits July 28, 2026 13:39

chore(nix-builder): remove empty dependencies from libtpu

fe2c742

tengomucho deployed to testpypi July 28, 2026 15:34 — with GitHub Actions View deployment

tengomucho requested a review from danieldk July 28, 2026 15:47

		@@ -0,0 +1,35 @@
		#!/usr/bin/env bash
		# Print the wheel URL and SRI hash for a torch-tpu-virtual-registry


		pub python_depends: Option<Vec<String>>,

		pub tpu: Option<TpuGeneral>,

Uh oh!

Conversation

tengomucho commented Jul 13, 2026

Important: Read before submitting

Checklist

Uh oh!

HuggingFaceDocBuilderDev commented Jul 13, 2026

Uh oh!

github-actions Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Coverage report — kernels/

Uh oh!

drbh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

danieldk Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danieldk left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

danieldk Jul 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

github-actions Bot commented Jul 13, 2026 •

edited

Loading

Coverage report — `kernels/`

danieldk Jul 27, 2026 •

edited

Loading

danieldk Jul 28, 2026 •

edited

Loading

tengomucho Jul 28, 2026 •

edited

Loading