Merge pull request #3457 from AI-Hypercomputer:BirdsOfAFthr-patch-1

PiperOrigin-RevId: 887044942

Author: Google-ML-Automation

src/maxtext/layers/moe.py

@@ -934,7 +934,7 @@ def gmm(
           output = mblx.gmm(
               lhs=inputs,
               rhs=kernel,
-              group_sizes=tokamax_group_sizes,
+              group_sizes=group_sizes,
               preferred_element_type=self.dtype,
               tiling=tiling,
               lhs_quantize_dtype=lhs_quantize_dtype,