PrivateUse1 backend
Also known as PrivateUse1, PyTorch out-of-tree backend, custom device backend
In plain English
PrivateUse1 is the PyTorch hook that lets a vendor plug in a new device type, so a tensor can live on device="tpu" just like it would on CUDA.
Technical definition
PrivateUse1 is the reserved dispatch key and device type in PyTorch that lets an out-of-tree backend register its own tensors, memory allocator, and operator implementations so the framework treats a new accelerator as a first-class device.
Engineering details
TorchTPU uses PrivateUse1 to expose an ordinary torch.Tensor on device="tpu" rather than a wrapper backed by a JAX array. The PyTorch dispatcher routes ATen operations to the TPU backend, which can execute eagerly for bring-up and debugging or capture graphs through torch.compile. Distributed APIs such as DDP, FSDP2, and DTensor work through the same mechanism. The earlier TorchAX approach instead intercepted operations through __torch_dispatch__ and translated each one into JAX calls.
Why it matters
The backend hook is what makes the word native concrete. Serving engines can reuse upstream model code, schedulers, continuous batching, and feature logic without a PyTorch-to-JAX translation layer, which lowers the cost of bringing each new model and engine feature to TPU.
How to read it in InferenceX
The TPU InferenceX preview results come from the TorchTPU vLLM stack built on this backend. Google expects to leave private beta and open source it around mid-October at the PyTorch Conference, after which SemiAnalysis moves TPU benchmarking from its fork to the public InferenceX repo.