Pins jax==0.4.23 (xla matched), flax==0.7.5, te==1.2.1

NVIDIA · Jan 23, 2024 · 489bacf · 489bacf
1 parent 5106983
commit 489bacf
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/.github/container/manifest.yaml b/.github/container/manifest.yaml
@@ -2,25 +2,25 @@
 jax:
   url: https://github.com/google/jax.git
   tracking_ref: main
-  latest_verified_commit: 2356e57af47baa37286177676a5a4bcd0c7b8c40
+  latest_verified_commit: jaxlib-v0.4.23
   mode: git-clone
 xla:
   url: https://github.com/openxla/xla.git
   tracking_ref: main
-  latest_verified_commit: 17ec7b4031608d3ddaca83e0661ed3aa0393e9d9
+  latest_verified_commit: a487d8ba5da8d0dec056972270e50de6748a1035
   mode: git-clone
 flax:
   url: https://github.com/google/flax.git
   mirror_url: https://github.com/nvjax-svc-0/flax.git
   tracking_ref: main
-  latest_verified_commit: 7cbc3c13ecd8a1713732113dec7939d0b75342c6
+  latest_verified_commit: v0.7.5
   mode: git-clone
   patches:
     pull/3340/head: file://patches/flax/PR-3340.patch # Add Sharding Annotations to Flax Modules
 transformer-engine:
   url: https://github.com/NVIDIA/TransformerEngine.git
   tracking_ref: main
-  latest_verified_commit: bb759adc2a53c57cc387f531c580a16149e3d4ac
+  latest_verified_commit: v1.2.1
   mode: git-clone
 t5x:
   url: https://github.com/google-research/t5x.git