Debug

abheesht17 · abheesht17 · commit 8a5cfcabf744 · 2025-10-27T16:05:53.000+05:30
diff --git a/examples/ml_perf/model.py b/examples/ml_perf/model.py
@@ -196,9 +196,11 @@ def call(self, inputs: dict[str, Tensor]) -> Tensor:
             for small_emb_feature in small_emb_inputs.keys():
                 small_emb_input = small_emb_inputs[small_emb_feature]
                 embedding_layer = self.small_embedding_layers[small_emb_feature]
+
                 embedding = embedding_layer(small_emb_input)
                 embedding = ops.sum(embedding, axis=-2)
                 small_embeddings.append(embedding)
+
             small_embeddings = ops.concatenate(small_embeddings, axis=-1)
 
         # Interaction
diff --git a/keras_rs/src/layers/embedding/jax/distributed_embedding.py b/keras_rs/src/layers/embedding/jax/distributed_embedding.py
@@ -412,6 +412,8 @@ def sparsecore_build(
         feature_specs = config_conversion.keras_to_jte_feature_configs(
             self._sc_feature_configs
         )
+        print(f"--->{self._sc_feature_configs=}")
+        print(f"--->{feature_specs=}")
 
         # Distribution for sparsecore operations.
         sparsecore_distribution, sparsecore_layout = (

Original file line number	Diff line number	Diff line change
`@@ -412,6 +412,8 @@ def sparsecore_build(`
`412`	`412`	`feature_specs = config_conversion.keras_to_jte_feature_configs(`
`413`	`413`	`self._sc_feature_configs`
`414`	`414`	`)`
	`415`	`+ print(f"--->{self._sc_feature_configs=}")`
	`416`	`+ print(f"--->{feature_specs=}")`
`415`	`417`
`416`	`418`	`# Distribution for sparsecore operations.`
`417`	`419`	`sparsecore_distribution, sparsecore_layout = (`