Allow customizing the stack name when using table_stacking.stack_tables.

PiperOrigin-RevId: 715449105

Author: Google-ML-Automation

jax_tpu_embedding/sparsecore/lib/nn/table_stacking.py

@@ -76,11 +76,6 @@ def round_up_dim_and_vocab_size(
   return table_to_padded_dim, tables_to_padded_vocab_size
 
 
-# TODO(b/355289256): Consider better name for the stack.
-def _get_stack_name(table_names: Sequence[str]) -> str:
-  return "_".join(table_names)
-
-
 def _get_stack_table_names(
     tables: Mapping[str, embedding_spec.TableSpec], num_sc: int
 ) -> Sequence[Sequence[str]]:
@@ -96,12 +91,13 @@ def _get_stack_table_names(
 
 
 def _verify_stack_tables(
+    stack_name: str,
     table_names: Sequence[str],
+    features: Sequence[embedding_spec.FeatureSpec],
     tables: Mapping[str, embedding_spec.TableSpec],
     table_to_padded_dim: Mapping[str, int],
 ):
   """Verifies that the provided stacking groups are valid."""
-  stack_name = _get_stack_name(table_names)
   logging.vlog(
       2,
       "Verifying stack group: %s with tables: %s",
@@ -122,6 +118,15 @@ def _is_stacked_already(table: embedding_spec.TableSpec):
           f" {tables[tname].setting_in_stack.stack_name}."
       )
 
+  for feature in features:
+    if (
+        _is_stacked_already(feature.table_spec)
+        and feature.table_spec.setting_in_stack.stack_name == stack_name
+    ):
+      raise ValueError(
+          f"Cannot use stack name {stack_name} since it's already used."
+      )
+
   # A table should not be repeated in a group.
   counter = collections.Counter(table_names)
   for table, count in counter.items():
@@ -214,6 +219,7 @@ def _get_limits_for_stack(
 
 
 def _stack_feature_specs(
+    stack_name: str,
     features: Nested[embedding_spec.FeatureSpec],
     table_names: Sequence[str],
     table_to_padded_dim: Mapping[str, int],
@@ -225,7 +231,6 @@ def _stack_feature_specs(
 ) -> None:
   """Updated the feature spec based on provided groups and stacking logic."""
 
-  stack_name = _get_stack_name(table_names)
   table_name_to_feature_spec = {
       f.table_spec.name: f for f in tree.flatten(features)
   }
@@ -324,6 +329,7 @@ def stack_tables(
     num_sc_per_device: int = 4,
     stack_to_max_ids_per_partition: LimitsCallable = get_default_limits,
     stack_to_max_unique_ids_per_partition: LimitsCallable = get_default_limits,
+    stack_name: str | None = None,
 ) -> None:
   """Creates new feature specs based on specified stacking groups.
 
@@ -341,7 +347,12 @@ def stack_tables(
       stack.
     stack_to_max_unique_ids_per_partition: Override the
       max_unique_ids_per_partition for each stack.
+    stack_name: A unique name for the table stack. If None, a default name will
+      be chosen.
   """
+  if stack_name is None:
+    # TODO(b/355289256): Consider better name for the stack.
+    stack_name = "_".join(table_names)
   flatten_features = tree.flatten(features)
   tables_in_group = {
       feature.table_spec.name: feature.table_spec
@@ -353,8 +364,15 @@ def stack_tables(
           tables_in_group, num_sc_per_device * global_device_count
       )
   )
-  _verify_stack_tables(table_names, tables_in_group, table_to_padded_dim)
+  _verify_stack_tables(
+      stack_name,
+      table_names,
+      flatten_features,
+      tables_in_group,
+      table_to_padded_dim,
+  )
   _stack_feature_specs(
+      stack_name=stack_name,
       features=features,
       table_names=table_names,
       table_to_padded_dim=table_to_padded_dim,

jax_tpu_embedding/sparsecore/lib/nn/tests/table_stacking_test.py

@@ -371,6 +371,82 @@ def limits(name: str, batch_size: int) -> int:
         0.5
     )
 
+  def test_manual_stacking_reuse_table_name(self):
+    table_spec_a = embedding_spec.TableSpec(
+        vocabulary_size=64,
+        embedding_dim=12,
+        initializer=lambda: jnp.zeros((64, 16), dtype=jnp.float32),
+        optimizer=embedding_spec.SGDOptimizerSpec(learning_rate=0.5),
+        combiner='sum',
+        name='A',
+        max_ids_per_partition=16,
+        max_unique_ids_per_partition=16,
+    )
+    table_spec_b = embedding_spec.TableSpec(
+        vocabulary_size=120,
+        embedding_dim=10,
+        initializer=lambda: jnp.zeros((128, 16), dtype=jnp.float32),
+        optimizer=embedding_spec.SGDOptimizerSpec(learning_rate=0.5),
+        combiner='sum',
+        name='B',
+        max_ids_per_partition=16,
+        max_unique_ids_per_partition=16,
+    )
+    table_spec_c = embedding_spec.TableSpec(
+        vocabulary_size=120,
+        embedding_dim=10,
+        initializer=lambda: jnp.zeros((128, 16), dtype=jnp.float32),
+        optimizer=embedding_spec.SGDOptimizerSpec(learning_rate=0.5),
+        combiner='sum',
+        name='C',
+        max_ids_per_partition=16,
+        max_unique_ids_per_partition=16,
+    )
+    feature_specs = [
+        embedding_spec.FeatureSpec(
+            table_spec=table_spec_a,
+            input_shape=(16, 1),
+            output_shape=(
+                16,
+                table_spec_a.embedding_dim,
+            ),
+            name='feature_a',
+        ),
+        embedding_spec.FeatureSpec(
+            table_spec=table_spec_b,
+            input_shape=(16, 1),
+            output_shape=(
+                16,
+                table_spec_b.embedding_dim,
+            ),
+            name='feature_b',
+        ),
+        embedding_spec.FeatureSpec(
+            table_spec=table_spec_c,
+            input_shape=(16, 1),
+            output_shape=(
+                16,
+                table_spec_c.embedding_dim,
+            ),
+            name='feature_c',
+        ),
+    ]
+
+    table_stacking.stack_tables(
+        feature_specs,
+        ('A', 'B'),
+        global_device_count=jax.device_count(),
+        stack_name='custom_stack',
+    )
+
+    with self.assertRaisesRegex(ValueError, 'custom_stack.*already used.*'):
+      table_stacking.stack_tables(
+          feature_specs,
+          ('C',),
+          global_device_count=jax.device_count(),
+          stack_name='custom_stack',
+      )
+
   def test_manual_stacking_not_same_optimizer(self):
     table_spec_a = embedding_spec.TableSpec(
         vocabulary_size=64,