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
4
Watchers
37
Size
3.7 MB
PythonApache License 2.0
deep-learningembeddingsjaxjax-scsparse-computationtpuxla
Created: Nov 30, 2022
Updated: Aug 24, 2026
Last push: Aug 24, 2026