address some comments

Signed-off-by: bzgoogle <beinuoz_google_com@t1v-n-3f245f84-w-0.us-central1-c.c.tpu-prod-env-one-vm.internal>

Author: bzgoogle

tests/models/jax/common/moe/test_deepseek_moe.py

@@ -1,8 +1,6 @@
 import os
 import unittest
 
-os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=8'
-
 import jax
 import jax.numpy as jnp
 import numpy as np

tpu_commons/models/jax/common/moe/deepseek_moe.py

@@ -325,14 +325,6 @@ def _unpermute(self, processed_tokens: jax.Array, sort_indices: jax.Array,
                 processed_tokens, jnp.argsort(sort_indices))
             reshaped_tokens_TXD = unsorted_tokens_tD.reshape(
                 -1, self.num_experts_per_tok, self.hidden_size)
-        # jax.debug.print(
-        #     "✅ reshaped_tokens_TXD on device:  reshaped_tokens_TXD[5]={t}",
-        #     t=reshaped_tokens_TXD[5, 0,:5]
-        # )
-        # jax.debug.print(
-        #     "✅ router_weights_TX  on device:  router_weights_TX={t}",
-        #     t=router_weights_TX[5, :]
-        # )
         with jax.named_scope("combine_weights"):
             output_TD = jnp.einsum(
                 "TXD,TX -> TD",
@@ -484,9 +476,6 @@ def _distributed_sparse_moe_fwd(
             compute_expert_ids = global_sorted_experts
             local_sorted_indices = jnp.arange(sorted_inputs.shape[0])
 
-        #debug_position_in_sorted = jnp.argsort(global_sort_indices)[40:48]
-        #debug_position_compute_inputs = jnp.argsort(local_sorted_indices)[debug_position_in_sorted]
-
         # 4. Compute: Apply experts using Grouped Matrix Multiply
         with jax.named_scope("gating"):
             # compute_inputs: (local total assignments, D)