Fix breakage in sparsecore preprocessor.

Hilly12 · recml authors · commit b4ade1fd3ab1 · 2025-07-21T17:19:30.000-07:00
PiperOrigin-RevId: 785632442
diff --git a/recml/layers/linen/sparsecore.py b/recml/layers/linen/sparsecore.py
@@ -339,6 +339,7 @@ def _to_np(x: Any) -> np.ndarray:
         global_device_count=self.sparsecore_config.global_device_count,
         num_sc_per_device=self.sparsecore_config.num_sc_per_device,
         sharding_strategy=self.sparsecore_config.sharding_strategy,
+        feature_stacking_strategy=embedding.FeatureStackingStrategy.STACK_THEN_SPLIT,
         allow_id_dropping=False,
         batch_number=self._batch_number,
     )

Original file line number	Diff line number	Diff line change
`@@ -339,6 +339,7 @@ def _to_np(x: Any) -> np.ndarray:`
`339`	`339`	`global_device_count=self.sparsecore_config.global_device_count,`
`340`	`340`	`num_sc_per_device=self.sparsecore_config.num_sc_per_device,`
`341`	`341`	`sharding_strategy=self.sparsecore_config.sharding_strategy,`
	`342`	`+ feature_stacking_strategy=embedding.FeatureStackingStrategy.STACK_THEN_SPLIT,`
`342`	`343`	`allow_id_dropping=False,`
`343`	`344`	`batch_number=self._batch_number,`
`344`	`345`	`)`