Use old code for preprocessing

Author: abheesht17

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

@@ -718,7 +718,7 @@ def _sparsecore_set_tables(self, tables: Mapping[str, ArrayLike]) -> None:
 
         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)
         sharded_tables = embedding_utils.stack_and_shard_tables(
             table_specs,
             tables,