Add DistributedEmbedding example for TPU on JAX. #2132
Conversation
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request adds a new example demonstrating the use of DistributedEmbedding with the JAX backend on TPUs. The example is provided as a Python script, a Jupyter notebook, and a markdown file for documentation. The changes are a valuable addition. My review focuses on improving the clarity of the documentation and fixing a few potential correctness issues in the example code. The main points are correcting misleading references to TensorFlow in a JAX example, fixing incorrect batch sizes in the distributed configuration, and ensuring the model is in the correct mode during evaluation.
hertschuh
force-pushed
the
jax_tpu_embedding
branch
2 times, most recently
from
3ae685d to
aeab875
Compare
This was run on a cloud TPU v6e-1.
hertschuh
force-pushed
the
jax_tpu_embedding
branch
from
aeab875 to
b8d7564
Compare
| input_shape=(BATCH_SIZE,), | ||
| output_shape=(BATCH_SIZE, EMBEDDING_DIMENSION), |
There was a problem hiding this comment.
Oh, one question. Is this BATCH_SIZE or PER_REPLICA_BATCH_SIZE? The documentation says PER_REPLICA_BATCH_SIZE: https://keras.io/keras_rs/api/embedding_layers/distributed_embedding/#configuration
There was a problem hiding this comment.
I will update the documentation. In the end, it's the global batch size.
4 participants
This was run on a cloud TPU v6e-1.