Merge pull request #2528 from AI-Hypercomputer:chengnuojin-mixtral-migration

Google-ML-Automation · Google-ML-Automation · commit 1c138f71090d · 2025-10-20T16:03:15.000-07:00
PiperOrigin-RevId: 821842379
diff --git a/src/MaxText/layers/mixtral.py b/src/MaxText/layers/mixtral.py
@@ -135,6 +135,8 @@ def __call__(
       deterministic,
       model_mode,
       previous_chunk=None,
+      page_state=None,
+      slot=None,
   ):
 
     inputs = nn.with_logical_constraint(inputs, self.activation_axis_names)
