absl-py
aqtp
array-record
cloud-accelerator-diagnostics
cloud-tpu-diagnostics
datasets
drjax>=0.1.4
flax
gcsfs
google-api-python-client
google-cloud-aiplatform
google-cloud-mldiagnostics>=0.5.10
google-cloud-monitoring
grain[parquet]
huggingface_hub
jax!=0.7.1
jaxlib!=0.7.1
jaxtyping
jsonlines
ml-collections
ml-goodput-measurement
numpy
omegaconf
optax
orbax-checkpoint
pathwaysutils
pillow
pre-commit
protobuf
pydantic
pyink
pylint
pytest
pytype
qwix>=0.1.6
sentencepiece
tensorboard-plugin-profile
tensorboardx
tensorflow-datasets
tensorflow-text
tensorflow
tiktoken
tokamax>=0.0.4
transformers
google-jetstream @ https://github.com/AI-Hypercomputer/JetStream/archive/29329e8e73820993f77cfc8efe34eb2a73f5de98.zip
mlperf-logging @ https://github.com/mlcommons/logging/archive/38ab22670527888c8eb7825a4ece176fcc36a95d.zip
