fix missing page_state

Author: NuojCheng

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)