diff --git a/praxis/pip_package/cloudbuild-release.yaml b/praxis/pip_package/cloudbuild-release.yaml index 1250b15f..1acb96ca 100644 --- a/praxis/pip_package/cloudbuild-release.yaml +++ b/praxis/pip_package/cloudbuild-release.yaml @@ -22,6 +22,7 @@ substitutions: options: dynamic_substitutions: true substitution_option: 'ALLOW_LOOSE' + machineType: E2_HIGHCPU_8 timeout: 5400s artifacts: objects: diff --git a/praxis/pip_package/requirements.in b/praxis/pip_package/requirements.in index a833a994..4bd1d3c5 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.15 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 6671b3a7..93d2ad12 100644 --- a/praxis/pip_package/requirements.txt +++ b/praxis/pip_package/requirements.txt @@ -83,7 +83,7 @@ executing==1.2.0 # via stack-data fastjsonschema==2.16.2 # via nbformat -fiddle @ git+https://github.com/google/fiddle +fiddle==0.2.4 # via -r praxis-requirements.in flatbuffers==1.12 # via tensorflow