Add TPU backend support (torch_tpu / Pallas kernels) - #714

Merged

tengomucho merged 34 commits into

mainfrom

Aug 3, 2026

Commits on Jul 13, 2026

kernels: add support for TPU
tengomucho
committed
9bd1937
kernels-data: add Backend::Tpu enum variant
tengomucho
committed
a780b4b
kernel-builder: detect TPU in noarch _ops.py template
tengomucho
committed
dac23cd
examples: add kernels-test/relu-tpu example kernel repository
tengomucho
committed
31c394b
kernels/pytest: register tpu_only marker
tengomucho
committed
36a9d72
nix-builder: add TPU backend support
tengomucho
committed
4359655
nix-builder: add TPU build configuration to versions.nix

tengomucho
and
claude
committed
ef27628
examples: add flake.nix to relu-tpu kernel

tengomucho
and
claude
committed
38712a1
nix-builder: build TPU wheels against the default python

tengomucho
and
claude
committed
ee10f69
kernels-data,kernels: resolve TPU python deps to nix packages

tengomucho
and
claude
committed
3cd1ac2
kernels,kernel-builder: detect TPU via torch.backends.tpu

tengomucho
and
claude
committed
cae1a7d
kernels-data: include TPU python-depends in backend metadata

tengomucho
and
claude
committed
4bf3946
examples/relu-tpu: annotate the jax function signature

tengomucho
and
claude
committed
7126d26

Commits on Jul 27, 2026

nix-builder: drop get_torch_tpu_hash.sh helper script
tengomucho
committed
68fe98d
Merge branch 'main' into tpu
tengomucho
committed
38a0df3
kernel-builder: drop tpu from init --backends help text
tengomucho
committed
ad70735
kernel-builder: drop neuron from init --backends help text
tengomucho
committed
226e1e4
docs: regenerate builder-cli.md for init --backends help text
tengomucho
committed
6f1e3a5
nix-builder: exclude tpu from the shared build cache

tengomucho
and
claude
committed
fb05cac
nix-builder: fetch libtpu/torch_tpu via requireFile, not fetchurl

tengomucho
and
claude
committed
da257e5
nix-builder: regenerate build-variants.json without tpu
tengomucho
committed
3de0b0d
docs: regenerate build-variants.md (tpu excluded from bundle cache)

tengomucho
and
claude
committed
4902e14

Commits on Jul 28, 2026

Commits on Jul 29, 2026

Merge branch 'main' into tpu
tengomucho
committed
0cfc585

Commits on Jul 30, 2026

python3Packages.torch-tpu: make import check work in the sandbox
danieldk
committed
ab33454

Commits on Jul 31, 2026

python3Packages.torch: fold torch-tpu into torch
danieldk
committed
286a048
python3Packages.libtpu: 0.0.43 -> 0.0.44.1
danieldk
committed
59dd5d4
Makefile: fixup `all` target
danieldk
committed
6e70381
Remove torch-tpu and torch-tpu kernel dependency
danieldk
committed
68676ad
Make `get-kernel-check` hook work on TPU
danieldk
committed
9864d91
nix-builder: make /sys faking in hooks conditional on `tpuSupport`
danieldk
committed
40427a7