Change heuristic based parameter static_buffer_size_multiplier to FDO tunable parameter feature_spec.table_spec.[stacked_table_spec.]suggested_coo_buffer_size.

PiperOrigin-RevId: 761725839

Author: adityagupta1089

recml/layers/linen/sparsecore.py

@@ -30,6 +30,7 @@
 from recml.core.ops import embedding_ops
 import tensorflow as tf
 
+
 with epy.lazy_imports():
   # pylint: disable=g-import-not-at-top
   from jax_tpu_embedding.sparsecore.lib.nn import embedding
@@ -99,8 +100,6 @@ class SparsecoreEmbedder:
       fixed mapping is used to determine this based on device 0. This may fail
       on newer TPU architectures if the mapping is not updated of if device 0 is
       not a TPU device with a sparsecore.
-    static_buffer_size_multiplier: The multiplier to use for the static buffer
-      size. Defaults to 256.
 
   Example usage:
   ```python
@@ -142,7 +141,6 @@ def __call__(self, inputs: Mapping[str, jax.Array]) -> jax.Array:
   specs: Mapping[str, EmbeddingSpec]
   optimizer: OptimizerSpec
   sharding_strategy: str = 'MOD'
-  static_buffer_size_multiplier: int = 256
 
   def __post_init__(self):
     self._feature_specs = None
@@ -254,7 +252,6 @@ def _preprocessor(inputs):
           global_device_count=jax.device_count(),
           num_sc_per_device=self._num_sc_per_device,
           sharding_strategy=self.sharding_strategy,
-          static_buffer_size_multiplier=self.static_buffer_size_multiplier,
           allow_id_dropping=False,
       )