diff --git a/praxis/pip_package/requirements.in b/praxis/pip_package/requirements.in index b46aa7a0..d6a36df7 100644 --- a/praxis/pip_package/requirements.in +++ b/praxis/pip_package/requirements.in @@ -4,7 +4,7 @@ absl-py clu einops -fiddle @ git+https://github.com/google/fiddle +fiddle flax jax jax-bitempered-loss @@ -15,4 +15,4 @@ optax-shampoo protobuf==3.19.5 t5x tensorflow~=2.9.2 -tensorflow-text~=2.9.0 \ No newline at end of file +tensorflow-text~=2.9.0 diff --git a/praxis/pip_package/requirements.txt b/praxis/pip_package/requirements.txt index 97247845..1c7ddd3e 100644 --- a/praxis/pip_package/requirements.txt +++ b/praxis/pip_package/requirements.txt @@ -89,7 +89,7 @@ executing==1.2.0 # via stack-data fastjsonschema==2.16.2 # via nbformat -fiddle @ git+https://github.com/google/fiddle +fiddle==0.2.5 # via -r praxis-requirements.in flatbuffers==1.12 # via tensorflow