JAX TPU Embedding (JAX SparseCore)

Quick Overview | Install guide | Reference docs

What is JAX TPU Embedding?

JAX SparseCore provides support for leveraging the SparseCore accelerators present in TPU generations starting with TPU v5p.

SparseCores are specialized tiled processors engineered for high-performance acceleration of workloads that involve irregular, sparse memory access and computation. While they excel at tasks like embedding lookups (common in deep learning recommendation models), their capabilities extend to accelerating a variety of other dynamic and sparse workloads on large datasets stored in High Bandwidth Memory (HBM).

This is a research project, not an official Google product. Expect sharp edges. Please help by trying it out, reporting bugs, and letting us know what you think!

Installation

You can install JAX TPU Embedding from PyPI:

pip install jax-tpu-embedding

Note: To use TPU acceleration, you must run in an environment with access to TPU v5p+ hardware and have the appropriate jax and jaxlib TPU releases installed (see the JAX installation guide).

Documentation

For detailed guides, specifications, and tutorials, see the JAX TPU Embedding Documentation.

• Quick Overview for Users
• Embedding Specification
• Input Processing
• Tutorials (Flax & Primitives)