update flax patch

Author: ashors1

.github/container/manifest.yaml

@@ -139,12 +139,12 @@ grain:
   latest_verified_commit: fa79b9dea81ffb00555a6c2ae2898be4bdd5e564
   mode: git-clone
 mujoco-mpc:
-  url:  https://github.com/google-deepmind/mujoco_mpc.git
+  url: https://github.com/google-deepmind/mujoco_mpc.git
   tracking_ref: main
   latest_verified_commit: a8bdcc1d968addf0d334243ce9bd6dbb07e1d781
   mode: git-clone
 language-to-reward-2023:
-  url:  https://github.com/google-deepmind/language_to_reward_2023.git
+  url: https://github.com/google-deepmind/language_to_reward_2023.git
   tracking_ref: main
   latest_verified_commit: abb8e5125e4ecd0da378490b73448c05a694def5
   mode: git-clone

.github/container/patches/flax/mirror-patch-t5x-layers-rebase-0131.patch

@@ -1,7 +1,7 @@
 From a563a57ca081e87b537bf724e2e29b138a14bb3f Mon Sep 17 00:00:00 2001
 From: ashors1 <ashors@nvidia.com>
 Date: Fri, 2 Jun 2023 15:01:21 -0700
-Subject: [PATCH 1/2] add t5x sharding annotations to flax layers
+Subject: [PATCH 1/3] add t5x sharding annotations to flax layers
 
 ---
  flax/linen/attention.py     | 34 +++++++++++++++++++++++-------
@@ -387,7 +387,7 @@ index cec8f508..e815f955 100644
 From 4ddde2bc878d9ff841f34e3826bbd1bce705766d Mon Sep 17 00:00:00 2001
 From: Terry Kong <terrycurtiskong@gmail.com>
 Date: Mon, 2 Oct 2023 16:10:05 -0700
-Subject: [PATCH 2/2] Added ConvTranspose sharding annotations (#3)
+Subject: [PATCH 2/3] Added ConvTranspose sharding annotations (#3)
 
 Co-authored-by: sahilj <sahilj@nvidia.com>
 ---
@@ -446,3 +446,54 @@ index 4656abf9..187ab6f5 100644
 -- 
 2.25.1
 
+
+From 9e68f4758d8a87553f3989e0a014974813e12390 Mon Sep 17 00:00:00 2001
+From: ashors1 <ashors@nvidia.com>
+Date: Thu, 1 Feb 2024 09:54:25 -0800
+Subject: [PATCH 3/3] Add missing import
+
+---
+ flax/linen/attention.py     | 1 +
+ flax/linen/linear.py        | 1 +
+ flax/linen/normalization.py | 1 +
+ 3 files changed, 3 insertions(+)
+
+diff --git a/flax/linen/attention.py b/flax/linen/attention.py
+index 2827119f..517ff2dc 100644
+--- a/flax/linen/attention.py
++++ b/flax/linen/attention.py
+@@ -39,6 +39,7 @@ from flax.typing import (
+   Initializer,
+   PrecisionLike,
+   DotGeneralT,
++  Tuple,
+ )
+ 
+ def dot_product_attention_weights(
+diff --git a/flax/linen/linear.py b/flax/linen/linear.py
+index 187ab6f5..759406ed 100644
+--- a/flax/linen/linear.py
++++ b/flax/linen/linear.py
+@@ -47,6 +47,7 @@ from flax.typing import (
+   ConvGeneralDilatedT,
+   PaddingLike,
+   LaxPadding,
++  Tuple,
+ )
+ 
+ 
+diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
+index e815f955..9e59e2be 100644
+--- a/flax/linen/normalization.py
++++ b/flax/linen/normalization.py
+@@ -32,6 +32,7 @@ from flax.typing import (
+   Shape as Shape,
+   Initializer,
+   Axes,
++  Tuple,
+ )
+ 
+ field = dataclasses.field
+-- 
+2.25.1
+