praxis/requirements.in at main · rajatsen91/praxis

25 lines (24 loc) · 528 Bytes

# To update requirements.txt for praxis and paxml, run:

# cd ../../paxml/pip_package && bash ./compile_requirements.sh

absl-py

# enforce newer chex version to avoid deprecated jax syntax

chex>=0.1.85

clu @ git+https://github.com/google/CommonLoopUtils

einops

etils

fiddle

flax

jax

jax-bitempered-loss

jaxtyping

numpy<2

optax

optax-shampoo

opt-einsum

# ensure sentencepiece is compataible with protobuf 3.19.

tensorflow-datasets>=4.8.3

tensorflow-metadata>=1.12.0

tensorflow-text>=2.9.0

tensorflow>=2.9.2

tfds-nightly

typeguard

Provide feedback