flax: 0.10.6 -> main (#1615)

Pin was added in #1602, now MaxText should support new versions:
AI-Hypercomputer/maxtext#2068.

Author: olupton

.github/container/manifest.yaml

@@ -11,8 +11,7 @@ xla:
 flax:
   url: https://github.com/google/flax.git
   mirror_url: https://github.com/nvjax-svc-0/flax.git
-  tracking_ref: v0.10.6
-  latest_verified_commit: 718aa8ccb12c3fdefcf3d196874e4fc667b3ade5
+  tracking_ref: main
   mode: git-clone
   patches:
     pull/3340/head: file://patches/flax/PR-3340.patch # Add Sharding Annotations to Flax Modules