requirements_lock.txt

# This file was autogenerated by uv via the following command:
#    bazel run //:requirements.update
absl-py==2.3.1
    # via
    #   -r requirements.in
    #   aqtp
    #   array-record
    #   chex
    #   clu
    #   dm-tree
    #   einshape
    #   etils
    #   grain
    #   keras
    #   ml-collections
    #   optax
    #   orbax-checkpoint
    #   seqio
    #   tensorboard
    #   tensorflow
    #   tensorflow-datasets
    #   tensorflow-metadata
    #   tokamax
aiofiles==25.1.0
    # via orbax-checkpoint
aiohappyeyeballs==2.6.1
    # via aiohttp
aiohttp==3.13.3
    # via
    #   fsspec
    #   gcsfs
aiosignal==1.4.0
    # via aiohttp
annotated-types==0.7.0
    # via pydantic
anyio==4.13.0
    # via
    #   google-genai
    #   httpx
aqtp==0.9.0
    # via -r requirements.in
array-record==0.8.3
    # via
    #   -r requirements.in
    #   grain
    #   tensorflow-datasets
astunparse==1.6.3
    # via tensorflow
attrs==25.4.0
    # via
    #   -r requirements.in
    #   aiohttp
    #   dm-tree
    #   hypothesis
cachetools==6.2.1
    # via google-auth
certifi==2025.10.5
    # via
    #   httpcore
    #   httpx
    #   kubernetes
    #   requests
    #   sentry-sdk
charset-normalizer==3.4.4
    # via requests
cheroot==11.1.2
    # via xprof
chex==0.1.90
    # via
    #   -r requirements.in
    #   optax
click==8.3.0
    # via
    #   nltk
    #   wandb
cloudpickle==3.1.1
    # via grain
clu==0.0.12
    # via seqio
datasets==4.8.4
    # via evaluate
decorator==5.2.1
    # via gcsfs
dill==0.4.1
    # via
    #   datasets
    #   evaluate
    #   multiprocess
distro==1.9.0
    # via google-genai
dm-tree==0.1.9
    # via tensorflow-datasets
docstring-parser==0.17.0
    # via
    #   google-cloud-aiplatform
    #   pyglove
    #   simple-parsing
durationpy==0.10
    # via kubernetes
editdistance==0.8.1
    # via seqio
einops==0.8.1
    # via
    #   -r requirements.in
    #   etils
einshape==1.0
    # via tokamax
etils[array-types, edc, enp, epath, epy, etqdm, etree]==1.13.0
    # via
    #   -r requirements.in
    #   array-record
    #   clu
    #   grain
    #   orbax-checkpoint
    #   tensorflow-datasets
    #   xprof
evaluate==0.4.6
    # via -r requirements.in
filelock==3.20.0
    # via
    #   datasets
    #   huggingface-hub
flatbuffers==25.9.23
    # via tensorflow
flax==0.11.2
    # via
    #   aqtp
    #   clu
    #   qwix
frozenlist==1.8.0
    # via
    #   aiohttp
    #   aiosignal
fsspec[http]==2025.10.0
    # via
    #   datasets
    #   etils
    #   evaluate
    #   gcsfs
    #   huggingface-hub
    #   xprof
gast==0.6.0
    # via tensorflow
gcsfs==2025.10.0
    # via xprof
gitdb==4.0.12
    # via gitpython
gitpython==3.1.45
    # via wandb
google-api-core[grpc]==2.28.1
    # via
    #   -r requirements.in
    #   google-api-python-client
    #   google-cloud-aiplatform
    #   google-cloud-appengine-logging
    #   google-cloud-bigquery
    #   google-cloud-build
    #   google-cloud-compute
    #   google-cloud-core
    #   google-cloud-logging
    #   google-cloud-monitoring
    #   google-cloud-resource-manager
    #   google-cloud-storage
    #   ml-goodput-measurement
google-api-python-client==2.194.0
    # via -r requirements.in
google-auth[requests]==2.42.1
    # via
    #   gcsfs
    #   google-api-core
    #   google-api-python-client
    #   google-auth-httplib2
    #   google-auth-oauthlib
    #   google-cloud-aiplatform
    #   google-cloud-appengine-logging
    #   google-cloud-bigquery
    #   google-cloud-build
    #   google-cloud-compute
    #   google-cloud-core
    #   google-cloud-logging
    #   google-cloud-monitoring
    #   google-cloud-resource-manager
    #   google-cloud-storage
    #   google-genai
google-auth-httplib2==0.3.1
    # via google-api-python-client
google-auth-oauthlib==1.2.4
    # via gcsfs
google-cloud-aiplatform==1.130.0
    # via -r requirements.in
google-cloud-appengine-logging==1.7.0
    # via google-cloud-logging
google-cloud-audit-log==0.4.0
    # via google-cloud-logging
google-cloud-bigquery==3.41.0
    # via google-cloud-aiplatform
google-cloud-build==3.36.0
    # via -r requirements.in
google-cloud-compute==1.47.0
    # via -r requirements.in
google-cloud-core==2.5.0
    # via
    #   google-cloud-bigquery
    #   google-cloud-logging
    #   google-cloud-storage
google-cloud-logging==3.12.1
    # via ml-goodput-measurement
google-cloud-monitoring==2.28.0
    # via ml-goodput-measurement
google-cloud-resource-manager==1.17.0
    # via google-cloud-aiplatform
google-cloud-storage==3.4.1
    # via
    #   -r requirements.in
    #   gcsfs
    #   google-cloud-aiplatform
google-crc32c==1.7.1
    # via
    #   google-cloud-storage
    #   google-resumable-media
google-genai==1.55.0
    # via google-cloud-aiplatform
google-pasta==0.2.0
    # via tensorflow
google-resumable-media==2.7.2
    # via
    #   google-cloud-bigquery
    #   google-cloud-storage
googleapis-common-protos[grpc]==1.71.0
    # via
    #   google-api-core
    #   google-cloud-audit-log
    #   grpc-google-iam-v1
    #   grpcio-status
    #   tensorflow-metadata
grain==0.2.14
    # via -r requirements.in
grpc-google-iam-v1==0.14.3
    # via
    #   google-cloud-build
    #   google-cloud-logging
    #   google-cloud-resource-manager
grpcio==1.76.0
    # via
    #   google-api-core
    #   google-cloud-appengine-logging
    #   google-cloud-build
    #   google-cloud-compute
    #   google-cloud-monitoring
    #   google-cloud-resource-manager
    #   googleapis-common-protos
    #   grpc-google-iam-v1
    #   grpcio-status
    #   tensorboard
    #   tensorflow
    #   tpu-info
grpcio-status==1.71.2
    # via google-api-core
gviz-api==1.10.0
    # via xprof
h11==0.16.0
    # via httpcore
h5py==3.15.1
    # via
    #   keras
    #   tensorflow
hf-xet==1.2.0
    # via huggingface-hub
httpcore==1.0.9
    # via httpx
httplib2==0.31.2
    # via
    #   google-api-python-client
    #   google-auth-httplib2
httpx==0.28.1
    # via
    #   datasets
    #   google-genai
huggingface-hub==0.36.0
    # via
    #   datasets
    #   evaluate
humanize==4.14.0
    # via orbax-checkpoint
hypothesis==6.143.0
    # via -r requirements.in
idna==3.11
    # via
    #   anyio
    #   httpx
    #   requests
    #   yarl
immutabledict==4.2.2
    # via
    #   tensorflow-datasets
    #   tokamax
importlib-metadata==8.7.0
    # via opentelemetry-api
importlib-resources==6.5.2
    # via etils
iniconfig==2.3.0
    # via pytest
jaraco-functools==4.4.0
    # via cheroot
jax==0.8.3
    # via
    #   -r requirements.in
    #   aqtp
    #   chex
    #   clu
    #   flax
    #   optax
    #   orbax-checkpoint
    #   qwix
    #   seqio
    #   tokamax
jaxlib==0.8.3
    # via
    #   -r requirements.in
    #   aqtp
    #   chex
    #   clu
    #   jax
    #   optax
    #   qwix
    #   seqio
    #   tokamax
jaxtyping==0.3.7
    # via tokamax
joblib==1.5.2
    # via
    #   nltk
    #   scikit-learn
keras==3.12.0
    # via tensorflow
kubernetes==35.0.0
    # via -r requirements.in
levenshtein==0.27.3
    # via -r requirements.in
libclang==18.1.1
    # via tensorflow
markdown==3.9
    # via tensorboard
markdown-it-py==4.0.0
    # via rich
markupsafe==3.0.3
    # via werkzeug
mdurl==0.1.2
    # via markdown-it-py
ml-collections==1.1.0
    # via clu
ml-dtypes==0.5.3
    # via
    #   jax
    #   jaxlib
    #   keras
    #   tensorflow
    #   tensorstore
ml-goodput-measurement==0.2.0
    # via -r requirements.in
more-itertools==10.8.0
    # via
    #   cheroot
    #   jaraco-functools
msgpack==1.1.2
    # via
    #   -r requirements.in
    #   flax
    #   orbax-checkpoint
multidict==6.7.1
    # via
    #   aiohttp
    #   yarl
multiprocess==0.70.19
    # via
    #   datasets
    #   evaluate
namex==0.1.0
    # via keras
nltk==3.9.2
    # via -r requirements.in
numpy==2.1.3
    # via
    #   -r requirements.in
    #   chex
    #   clu
    #   datasets
    #   dm-tree
    #   einshape
    #   etils
    #   evaluate
    #   flax
    #   grain
    #   h5py
    #   jax
    #   jaxlib
    #   keras
    #   ml-dtypes
    #   ml-goodput-measurement
    #   optax
    #   orbax-checkpoint
    #   pandas
    #   pycocotools
    #   scikit-learn
    #   scipy
    #   seqio
    #   shapely
    #   tensorboard
    #   tensorboardx
    #   tensorflow
    #   tensorflow-datasets
    #   tensorstore
    #   treescope
oauthlib==3.3.1
    # via requests-oauthlib
opentelemetry-api==1.38.0
    # via google-cloud-logging
opt-einsum==3.4.0
    # via
    #   jax
    #   tensorflow
optax==0.2.6
    # via
    #   -r requirements.in
    #   flax
optree==0.17.0
    # via keras
orbax-checkpoint==0.11.33
    # via
    #   -r requirements.in
    #   flax
packaging==25.0
    # via
    #   clu
    #   datasets
    #   evaluate
    #   google-cloud-aiplatform
    #   google-cloud-bigquery
    #   huggingface-hub
    #   keras
    #   pytest
    #   seqio
    #   tensorboard
    #   tensorboardx
    #   tensorflow
    #   tpu-info
    #   wandb
pandas==3.0.2
    # via
    #   datasets
    #   evaluate
pika==1.3.2
    # via -r requirements.in
pillow==12.0.0
    # via -r requirements.in
platformdirs==4.5.0
    # via wandb
pluggy==1.6.0
    # via pytest
portpicker==1.6.0
    # via -r requirements.in
prefixed==0.9.0
    # via -r requirements.in
prometheus-client==0.25.0
    # via -r requirements.in
promise==2.3
    # via tensorflow-datasets
propcache==0.4.1
    # via
    #   aiohttp
    #   yarl
proto-plus==1.26.1
    # via
    #   google-api-core
    #   google-cloud-aiplatform
    #   google-cloud-appengine-logging
    #   google-cloud-build
    #   google-cloud-compute
    #   google-cloud-logging
    #   google-cloud-monitoring
    #   google-cloud-resource-manager
protobuf==5.29.5
    # via
    #   google-api-core
    #   google-cloud-aiplatform
    #   google-cloud-appengine-logging
    #   google-cloud-audit-log
    #   google-cloud-build
    #   google-cloud-compute
    #   google-cloud-logging
    #   google-cloud-monitoring
    #   google-cloud-resource-manager
    #   googleapis-common-protos
    #   grain
    #   grpc-google-iam-v1
    #   grpcio-status
    #   orbax-checkpoint
    #   proto-plus
    #   tensorboard
    #   tensorboardx
    #   tensorflow
    #   tensorflow-datasets
    #   tensorflow-metadata
    #   tpu-info
    #   wandb
    #   xprof
psutil==7.1.2
    # via
    #   -r requirements.in
    #   orbax-checkpoint
    #   portpicker
    #   tensorflow-datasets
pyarrow==22.0.0
    # via
    #   datasets
    #   tensorflow-datasets
pyasn1==0.6.1
    # via
    #   pyasn1-modules
    #   rsa
pyasn1-modules==0.4.2
    # via google-auth
pycocotools==2.0.11
    # via -r requirements.in
pydantic==2.12.4
    # via
    #   google-cloud-aiplatform
    #   google-genai
    #   tokamax
    #   wandb
pydantic-core==2.41.5
    # via pydantic
pyglove==0.4.5
    # via seqio
pygments==2.19.2
    # via
    #   pytest
    #   rich
pyparsing==3.2.5
    # via httplib2
pytest==8.4.2
    # via -r requirements.in
python-dateutil==2.9.0.post0
    # via
    #   google-cloud-bigquery
    #   kubernetes
    #   pandas
pyyaml==6.0.3
    # via
    #   datasets
    #   flax
    #   huggingface-hub
    #   kubernetes
    #   ml-collections
    #   orbax-checkpoint
    #   wandb
qwix==0.1.2
    # via tokamax
rapidfuzz==3.14.5
    # via levenshtein
regex==2025.10.23
    # via nltk
requests==2.32.5
    # via
    #   -r requirements.in
    #   datasets
    #   evaluate
    #   gcsfs
    #   google-api-core
    #   google-auth
    #   google-cloud-bigquery
    #   google-cloud-storage
    #   google-genai
    #   huggingface-hub
    #   kubernetes
    #   ml-goodput-measurement
    #   requests-oauthlib
    #   tensorflow
    #   tensorflow-datasets
    #   wandb
requests-oauthlib==2.0.0
    # via
    #   google-auth-oauthlib
    #   kubernetes
rich==14.2.0
    # via
    #   flax
    #   keras
    #   tpu-info
rsa==4.9.1
    # via google-auth
scikit-learn==1.6.1
    # via -r requirements.in
scipy==1.16.3
    # via
    #   -r requirements.in
    #   jax
    #   jaxlib
    #   ml-goodput-measurement
    #   scikit-learn
sentencepiece==0.2.1
    # via seqio
sentry-sdk==2.43.0
    # via wandb
seqio==0.0.20
    # via -r requirements.in
setuptools==80.9.0
    # via
    #   chex
    #   tensorboard
    #   tensorflow
    #   xprof
shapely==2.1.2
    # via google-cloud-aiplatform
simple-parsing==0.1.7
    # via tensorflow-datasets
simplejson==3.20.2
    # via orbax-checkpoint
six==1.17.0
    # via
    #   astunparse
    #   google-pasta
    #   gviz-api
    #   kubernetes
    #   promise
    #   python-dateutil
    #   tensorboard
    #   tensorflow
    #   xprof
smmap==5.0.2
    # via gitdb
sniffio==1.3.1
    # via google-genai
sortedcontainers==2.4.0
    # via hypothesis
tenacity==9.1.4
    # via google-genai
tensorboard==2.19.0
    # via tensorflow
tensorboard-data-server==0.7.2
    # via tensorboard
tensorboardx==2.6.4
    # via
    #   ml-goodput-measurement
    #   tokamax
tensorflow==2.19.1
    # via
    #   -r requirements.in
    #   tensorflow-text
tensorflow-datasets==4.9.9
    # via
    #   -r requirements.in
    #   seqio
tensorflow-io @ https://storage.googleapis.com/axlearn-wheels/tensorflow_io/tensorflow_io-0.37.3-cp312-cp312-manylinux_2_31_x86_64.whl
    # via -r requirements.in
tensorflow-io-gcs-filesystem @ https://storage.googleapis.com/axlearn-wheels/tensorflow_io_gcs_filesystem/tensorflow_io_gcs_filesystem-0.37.3-cp312-cp312-manylinux_2_31_x86_64.whl
    # via
    #   -r requirements.in
    #   tensorflow-io
tensorflow-metadata==1.17.2
    # via tensorflow-datasets
tensorflow-text==2.19.0
    # via seqio
tensorstore==0.1.78
    # via
    #   -r requirements.in
    #   flax
    #   orbax-checkpoint
termcolor==3.2.0
    # via
    #   pyglove
    #   tensorflow
    #   tensorflow-datasets
threadpoolctl==3.6.0
    # via scikit-learn
tokamax==0.0.10
    # via -r requirements.in
toml==0.10.2
    # via tensorflow-datasets
toolz==1.1.0
    # via chex
tpu-info==0.11.0
    # via -r requirements.in
tqdm==4.67.1
    # via
    #   -r requirements.in
    #   datasets
    #   etils
    #   evaluate
    #   huggingface-hub
    #   nltk
    #   tensorflow-datasets
    #   tokamax
treescope==0.1.10
    # via flax
typeguard==2.13.3
    # via tokamax
typing-extensions==4.15.0
    # via
    #   -r requirements.in
    #   aiosignal
    #   anyio
    #   chex
    #   clu
    #   etils
    #   flax
    #   google-cloud-aiplatform
    #   google-genai
    #   grpcio
    #   huggingface-hub
    #   opentelemetry-api
    #   optree
    #   orbax-checkpoint
    #   pydantic
    #   pydantic-core
    #   simple-parsing
    #   tensorflow
    #   tokamax
    #   typing-inspection
    #   wandb
typing-inspection==0.4.2
    # via pydantic
uritemplate==4.2.0
    # via google-api-python-client
urllib3==2.5.0
    # via
    #   kubernetes
    #   ml-goodput-measurement
    #   requests
    #   sentry-sdk
uvloop==0.22.1
    # via orbax-checkpoint
wadler-lindig==0.1.7
    # via jaxtyping
wandb==0.22.3
    # via -r requirements.in
websocket-client==1.9.0
    # via kubernetes
websockets==15.0.1
    # via google-genai
werkzeug==3.1.3
    # via
    #   tensorboard
    #   xprof
wheel==0.45.1
    # via astunparse
wrapt==2.0.0
    # via
    #   -r requirements.in
    #   clu
    #   dm-tree
    #   tensorflow
    #   tensorflow-datasets
xprof==2.21.5
    # via tokamax
xxhash==3.6.0
    # via
    #   datasets
    #   evaluate
yarl==1.22.0
    # via aiohttp
zipp==3.23.0
    # via
    #   etils
    #   importlib-metadata