Links for jax-tpu-embedding