jax-tpu-embedding
jax-ml/jax-tpu-embedding
High-performance large-scale embedding acceleration for JAX on Google TPU SparseCores.
37stars
Forks
5
Open issues
11
Watchers
37
Size
4.1 MB
PythonApache License 2.0
deep-learningembeddingsjaxjax-scsparse-computationtpuxla
Created: Nov 30, 2022
Updated: Oct 9, 2026
Last push: Oct 9, 2026