Merge pull request #3541 from AI-Hypercomputer:parambole/494634065

PiperOrigin-RevId: 895394473

Author: Google-ML-Automation

src/maxtext/layers/moe.py

@@ -1182,7 +1182,7 @@ def wrapper(x, logits, pre_bias_logits, w0, w1, wo, w0_bias, w1_bias, wo_bias, r
             # experts_per_shard > num_experts_per_tok we cannot assign more than num_experts_per_tok to all of the inputs.
             max_local_experts_per_tok = min(local_expert_size, self.config.num_experts_per_tok)
             buffer_size = int(num_expert_parallelism * batch_size * sequence_length * max_local_experts_per_tok)
-            output_shape = jnp.zeros((buffer_size, self.config.emb_dim), dtype=x.dtype)
+            output_shape = jax.lax.empty((buffer_size, self.config.emb_dim), dtype=x.dtype)
 
             x = jax.lax.ragged_all_to_all(
                 x,
@@ -1345,7 +1345,7 @@ def get_active_sharding_axes(pspec_dim_axes, tensor_dim_index):
           original_inputs_first_dim = batch_size * sequence_length * self.config.num_experts_per_tok
           if sorted_selected_experts.shape[0] != original_inputs_first_dim:
             raise ValueError("original_inputs_first_dim does not match the original tensor" " shape!")
-          output_shape = jnp.zeros(
+          output_shape = jax.lax.empty(
               (
                   original_inputs_first_dim,
                   self.config.emb_dim // self.get_tensor_parallelism_size(),