Add debugging statements

Author: abheesht17

examples/ml_perf/run.sh

@@ -144,6 +144,22 @@ gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
   --worker=all \
   --command="source .keras-env/bin/activate && pip install -U 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html"
 
+# ==============================================================================
+# Kill Previous Training Processes
+# ==============================================================================
+# echo ">>> Listing matching processes..."
+# gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
+#   --project ${PROJECT} \
+#   --zone ${ZONE} \
+#   --worker=all \
+#   --command="ps aux | grep '[e]xamples.ml_perf.main' || true"
+
+# echo ">>> Terminating any existing training processes..."
+# gcloud alpha compute tpus tpu-vm ssh ${TPU_NAME} \
+#   --project ${PROJECT} \
+#   --zone ${ZONE} \
+#   --worker=all \
+#   --command="pkill -9 -f 'python3.12 -m examples.ml_perf.[m]ain.*' || true"
 
 # ==============================================================================
 # Verify Installation

keras_rs/src/layers/embedding/jax/distributed_embedding.py

@@ -580,12 +580,17 @@ def _sparsecore_preprocess(
         )
 
         layout = self._sparsecore_layout
+        print(f"-->{layout=}")
         mesh = layout.device_mesh.backend_mesh
+        print(f"-->{mesh=}")
         global_device_count = mesh.devices.size
+        print(f"-->{global_device_count=}")
         local_device_count = mesh.local_mesh.devices.size
+        print(f"{local_device_count=}")
         num_sc_per_device = jte_utils.num_sparsecores_per_device(
             mesh.devices.item(0)
         )
+        print(f"-->{num_sc_per_device=}")
 
         preprocessed, stats = embedding_utils.stack_and_shard_samples(
             self._config.feature_specs,
@@ -594,13 +599,15 @@ def _sparsecore_preprocess(
             global_device_count,
             num_sc_per_device,
         )
+        print(f"-->{stats=}")
 
         if training:
             # Synchronize input statistics across all devices and update the
             # underlying stacked tables specs in the feature specs.
 
             # Aggregate stats across all processes/devices via pmax.
             num_local_cpu_devices = jax.local_device_count("cpu")
+            print(f"-->{num_local_cpu_devices=}")
 
             def pmax_aggregate(x: Any) -> Any:
                 if not hasattr(x, "ndim"):