Use old code for preprocessing

Author: abheesht17

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

@@ -441,7 +441,7 @@ def sparsecore_build(
         )
 
         # Collect all stacked tables.
-        table_specs = embedding_utils.get_table_specs(feature_specs)
+        table_specs = embedding.get_table_specs(feature_specs)
         table_stacks = embedding_utils.get_table_stacks(table_specs)
 
         # Create variables for all stacked tables and slot variables.
@@ -515,7 +515,7 @@ def _sparsecore_symbolic_preprocess(
         del inputs, weights, training
 
         # Each stacked-table gets a ShardedCooMatrix.
-        table_specs = embedding_utils.get_table_specs(
+        table_specs = embedding.get_table_specs(
             self._config.feature_specs
         )
         table_stacks = embedding_utils.get_table_stacks(table_specs)
@@ -750,7 +750,7 @@ def _sparsecore_get_embedding_tables(self) -> dict[str, ArrayLike]:
 
         config = self._config
         num_table_shards = config.mesh.devices.size * config.num_sc_per_device
-        table_specs = embedding_utils.get_table_specs(config.feature_specs)
+        table_specs = embedding.get_table_specs(config.feature_specs)
 
         # Extract only the table variables, not the gradient slot variables.
         table_variables = {