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 v5.
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 v5+ 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.
2 maintainers
Metadata
2 maintainers
Release files for jax-tpu-embedding 0.2.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| jax_tpu_embedding-0.2.0.tar.gz | 4.1 MB | Details |
Built distributions (wheels)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| jax_tpu_embedding-0.2.0-cp314-cp314-manylinux_2_31_x86_64.whl | CPython 3.14 | CPython 3.14 | Linux glibc 2.31+ x86-64 | Details |
| jax_tpu_embedding-0.2.0-cp313-cp313-manylinux_2_31_x86_64.whl | CPython 3.13 | CPython 3.13 | Linux glibc 2.31+ x86-64 | Details |
| jax_tpu_embedding-0.2.0-cp312-cp312-manylinux_2_31_x86_64.whl | CPython 3.12 | CPython 3.12 | Linux glibc 2.31+ x86-64 | Details |
Total release size: 16.6 MB
Release files / jax_tpu_embedding-0.2.0.tar.gz
| Download URL | jax_tpu_embedding-0.2.0.tar.gz |
|---|---|
| Size | 4.1 MB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
490f0ba077ff4f3b16486229d8383aa372348ef6601a5e952646fdcb5b322a99
|
|
BLAKE2b-256 checksum How to use checksums |
9a0a219af3d6292ddc51549ed08104f3522b8e81c01e175c4151a3d2928ddc40
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.12.13
|
Release files / jax_tpu_embedding-0.2.0-cp314-cp314-manylinux_2_31_x86_64.whl
| Download URL | jax_tpu_embedding-0.2.0-cp314-cp314-manylinux_2_31_x86_64.whl |
|---|---|
| Size | 4.2 MB |
| Tags | CPython 3.14 Linux glibc 2.31+ x86-64 |
|
SHA-256 checksum How to use checksums |
0b2a993a558244bbc2d1f2e3c47ae4d7c835914c6bcb67ce977e28e0bf3d9232
|
|
BLAKE2b-256 checksum How to use checksums |
956eea095750a01d72da6b26bf1ac3ea9cd1571d28a0d6f63a93fa835f63a253
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.12.13
|
Release files / jax_tpu_embedding-0.2.0-cp313-cp313-manylinux_2_31_x86_64.whl
| Download URL | jax_tpu_embedding-0.2.0-cp313-cp313-manylinux_2_31_x86_64.whl |
|---|---|
| Size | 4.2 MB |
| Tags | CPython 3.13 Linux glibc 2.31+ x86-64 |
|
SHA-256 checksum How to use checksums |
2f722093682d56e8f108c9cf811ecd38b8cca66355e6e5218b01b8123c6415c3
|
|
BLAKE2b-256 checksum How to use checksums |
c2b91135b54e92003410e18cc1f82becdec8d234318574c919ee8304666621b5
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.12.13
|
Release files / jax_tpu_embedding-0.2.0-cp312-cp312-manylinux_2_31_x86_64.whl
| Download URL | jax_tpu_embedding-0.2.0-cp312-cp312-manylinux_2_31_x86_64.whl |
|---|---|
| Size | 4.2 MB |
| Tags | CPython 3.12 Linux glibc 2.31+ x86-64 |
|
SHA-256 checksum How to use checksums |
b968d57a5e386e176bf79ed2b75957a8413671da6305c77e114756998e019f13
|
|
BLAKE2b-256 checksum How to use checksums |
a3e0b47d807d03c1727030c0b2f3ac5091144edc9e2589e66d43bdb832bc9ed9
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
Yes |
| Uploaded via |
twine/6.1.0 CPython/3.12.13
|
