[JAX SC] Raise error on updating FDO params after stacking.

PiperOrigin-RevId: 826192844

Author: adityagupta1089

jax_tpu_embedding/sparsecore/lib/core/pybind_input_preprocessing_test.py

@@ -2923,10 +2923,12 @@ def test_combiner(
     )
 
     # Compute expected by re-adjusting the weights using a "sum" combiner.
-    table_spec.combiner = "sum"
-    table_spec.stacked_table_spec = None
+    table_spec_for_expectation = dataclasses.replace(table_spec, combiner="sum")
+    feature_spec_for_expectation = dataclasses.replace(
+        feature_spec, table_spec=table_spec_for_expectation
+    )
     embedding.prepare_feature_specs_for_training(
-        feature_spec,
+        feature_spec_for_expectation,
         global_device_count=1,
         num_sc_per_device=4,
     )
@@ -2941,7 +2943,7 @@ def test_combiner(
     ) = pybind_input_preprocessing.PreprocessSparseDenseMatmulInput(
         [input_features],
         [input_weights],
-        [feature_spec],
+        [feature_spec_for_expectation],
         local_device_count=1,
         global_device_count=1,
         num_sc_per_device=4,

jax_tpu_embedding/sparsecore/lib/core/sparse_coo_input_preprocessing_test.py

@@ -101,7 +101,6 @@ def test_sparse_tensor_input(self, has_leading_dimension):
     indices_tensor = [sparse_tensor.indices]
     values_tensor = [sparse_tensor.values]
     dense_shape_tensor = [sparse_tensor.dense_shape]
-    self.feature_spec.table_spec.suggested_coo_buffer_size_per_device = 64
     batch_number = 42
     (
         row_pointers_sparse,
@@ -148,7 +147,6 @@ def test_sparse_tensor_input(self, has_leading_dimension):
           dtype=np.int32,
       )
 
-    self.feature_spec.table_spec.suggested_coo_buffer_size_per_device = 64
     batch_number = 42
     local_device_count = 4
     global_device_count = 4
@@ -246,7 +244,6 @@ def test_sparse_tensor_input_with_empty_rows(self, has_leading_dimension):
     numpy_features = np.array(numpy_features, dtype=object)
     numpy_weights = np.array(numpy_weights, dtype=object)
 
-    self.feature_spec.table_spec.suggested_coo_buffer_size_per_device = 64
     batch_number = 42
     local_device_count = 4
     global_device_count = 4
@@ -1374,7 +1371,6 @@ def test_sparse_tensor_input(self, has_leading_dimension):
     indices_tensor = [sparse_tensor.indices]
     values_tensor = [sparse_tensor.values]
     dense_shape_tensor = [sparse_tensor.dense_shape]
-    self.feature_spec.table_spec.suggested_coo_buffer_size_per_device = 64
     batch_number = 42
     (
         row_pointers_sparse,
@@ -1420,7 +1416,6 @@ def test_sparse_tensor_input(self, has_leading_dimension):
           ],
           dtype=np.int32,
       )
-    self.feature_spec.table_spec.suggested_coo_buffer_size_per_device = 64
     batch_number = 42
     local_device_count = 4
     global_device_count = 4

jax_tpu_embedding/sparsecore/lib/nn/embedding_spec.py

@@ -19,7 +19,7 @@
 import collections
 import dataclasses
 import inspect
-from typing import Callable, Sequence, TypeAlias
+from typing import Any, Callable, Sequence, TypeAlias
 
 from flax import struct
 import jax
@@ -673,6 +673,52 @@ class TableSpec:
   """Quantization config (min, max, num_buckets) which represent the float
   range and number of discrete integer buckets to use for quantization."""
 
+  _initialized: bool = dataclasses.field(
+      init=False, default=False, compare=False
+  )
+
+  def __setattr__(self, name: str, value: Any):
+    # These attributes are allowed to be modified after stacking because
+    # they are part of the stacking process itself (which sets
+    # `stacked_table_spec` and `setting_in_stack`), or are used to update
+    # stacking-related parameters like buffer sizes via replacing the
+    # StackedTableSpec.
+    allowed_after_stacking = (
+        "_stacked_table_spec",
+        "_setting_in_stack",
+        "stacked_table_spec",
+        "setting_in_stack",
+    )
+    fdo_params = (
+        "max_ids_per_partition",
+        "max_unique_ids_per_partition",
+        "suggested_coo_buffer_size_per_device",
+    )
+    # Check if __post_init__ has run and if this is a stacked table
+    if (
+        name not in allowed_after_stacking
+        and self._initialized
+        and self.is_stacked()
+    ):
+      if name in fdo_params:
+        raise AttributeError(
+            f"Cannot update parameter '{name}' on TableSpec '{self.name}' "
+            "after it has been stacked. If you are trying to update FDO "
+            "parameters like 'max_ids_per_partition', use "
+            "embedding.update_preprocessing_parameters() to update the "
+            "StackedTableSpec instead."
+        )
+      else:
+        raise AttributeError(
+            f"Cannot update parameter '{name}' on TableSpec '{self.name}' "
+            "after it has been stacked. If you need to modify non-FDO "
+            "parameters for a stacked table, create a new table spec using "
+            "`dataclasses.replace(table_spec, ...)` and feature spec using "
+            "`dataclasses.replace(feature_spec, table_spec=new_table_spec)` "
+            "and then re-prepare feature specs for training."
+        )
+    super().__setattr__(name, value)
+
   # This points to the stacked table spec which this table belongs to.
   # If this is None, this table is the top-most table.
   _stacked_table_spec: StackedTableSpec | None = dataclasses.field(
@@ -721,6 +767,7 @@ def __post_init__(
           row_offset_in_shard=0,
           shard_rotation=0,
       )
+    object.__setattr__(self, "_initialized", True)
 
 
 @dataclasses.dataclass(eq=True, unsafe_hash=True, kw_only=True)