Remove auto stack kwargs

Author: abheesht17

examples/ml_perf/main.py

@@ -117,12 +117,6 @@ def main(
         top_mlp_dims=model_cfg.top_mlp_dims,
         num_dcn_layers=model_cfg.num_dcn_layers,
         dcn_projection_dim=model_cfg.dcn_projection_dim,
-        auto_stack_kwargs={
-            "max_ids_per_partition": model_cfg.max_ids_per_partition,
-            "max_unique_ids_per_partition": (
-                model_cfg.max_unique_ids_per_partition
-            ),
-        },
         seed=SEED,
         dtype="float32",
         name="dlrm_dcn_v2",

examples/ml_perf/model.py

@@ -49,7 +49,6 @@ def __init__(
         top_mlp_dims: list[int],
         num_dcn_layers: int,
         dcn_projection_dim: int,
-        auto_stack_kwargs: dict[str, Any],
         seed: int | keras.random.SeedGenerator | None = None,
         dtype: str | None = None,
         name: str | None = None,
@@ -116,7 +115,6 @@ def __init__(
         self.embedding_layer = keras_rs.layers.DistributedEmbedding(
             feature_configs=large_emb_feature_configs,
             table_stacking="auto",
-            auto_stack_kwargs=auto_stack_kwargs,
             dtype=dtype,
             name="embedding_layer",
         )
@@ -173,7 +171,6 @@ def __init__(
         self.top_mlp_dims = top_mlp_dims
         self.num_dcn_layers = num_dcn_layers
         self.dcn_projection_dim = dcn_projection_dim
-        self.auto_stack_kwargs = auto_stack_kwargs
 
     def call(self, inputs: dict[str, Tensor]) -> Tensor:
         """Forward pass of the model.
@@ -280,7 +277,6 @@ def get_config(self):
                 "top_mlp_dims": self.top_mlp_dims,
                 "num_dcn_layers": self.num_dcn_layers,
                 "dcn_projection_dim": self.dcn_projection_dim,
-                "auto_stack_kwargs": self.auto_stack_kwargs,
                 "seed": self.seed,
             }
         )

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

@@ -193,39 +193,6 @@ def __call__(
 class DistributedEmbedding(base_distributed_embedding.DistributedEmbedding):
     """JAX implementation of the TPU embedding layer."""
 
-    def __init__(self, **kwargs: Any):
-        # Pull out `auto_stack_kwargs` from `kwargs`.
-        auto_stack_kwargs = kwargs.pop("auto_stack_kwargs", {})
-
-        auto_stack_max_ids_per_partition = auto_stack_kwargs.pop(
-            "max_ids_per_partition", None
-        )
-        auto_stack_max_unique_ids_per_partition = auto_stack_kwargs.pop(
-            "max_unique_ids_per_partition", None
-        )
-
-        # For `max_ids_per_partition` and `max_unique_ids_per_partition`, JTE's
-        # `auto_stack_tables` expects callables.
-        def _get_max_ids_per_partition(name: str, batch_size: int) -> int:
-            return auto_stack_max_ids_per_partition
-
-        def _get_max_unique_ids_per_partition(
-            name: str, batch_size: int
-        ) -> int:
-            return auto_stack_max_unique_ids_per_partition
-
-        if auto_stack_max_ids_per_partition is not None:
-            auto_stack_kwargs["stack_to_max_ids_per_partition"] = (
-                _get_max_ids_per_partition
-            )
-        if auto_stack_max_unique_ids_per_partition is not None:
-            auto_stack_kwargs["stack_to_max_unique_ids_per_partition"] = (
-                _get_max_unique_ids_per_partition
-            )
-
-        self._auto_stack_kwargs = auto_stack_kwargs
-        super().__init__(**kwargs)
-
     def _create_sparsecore_distribution(
         self, sparsecore_axis_name: str = "sparsecore"
     ) -> tuple[
@@ -438,7 +405,6 @@ def sparsecore_build(
                         feature_specs,
                         global_device_count,
                         num_sc_per_device,
-                        # **self._auto_stack_kwargs,
                     )
                 else:
                     raise ValueError(