Update moe.py

Author: suexu1025

src/MaxText/layers/moe.py

@@ -1673,7 +1673,7 @@ def dense_matmul(
         if self.config.activations_in_float32:
           intermediate_layer = intermediate_layer.astype(jnp.float32)
         intermediate_layer = adc.checkpoint_name(intermediate_layer, "mlpwo")
-      with jax.named_scope("w_sum"):
+      with jax.named_scope("weight_sum"):
         if is_llama4_decoder_layer:
           weights = self.reshape_and_update_weights(jnp.ones_like(top_k_weights), top_k_indices)
         if self.config.float32_weight_sum: