Use old code for preprocessing

Author: abheesht17

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

@@ -442,7 +442,7 @@ def sparsecore_build(
 
         # Collect all stacked tables.
         table_specs = embedding.get_table_specs(feature_specs)
-        table_stacks = embedding_utils.get_table_stacks(table_specs)
+        table_stacks = jte_table_stacking.get_table_stacks(table_specs)
 
         # Create variables for all stacked tables and slot variables.
         with sparsecore_distribution.scope():
@@ -518,7 +518,7 @@ def _sparsecore_symbolic_preprocess(
         table_specs = embedding.get_table_specs(
             self._config.feature_specs
         )
-        table_stacks = embedding_utils.get_table_stacks(table_specs)
+        table_stacks = jte_table_stacking.get_table_stacks(table_specs)
         stacked_table_specs = {
             stack_name: stack[0].stacked_table_spec
             for stack_name, stack in table_stacks.items()