Allow construction SDMM stats object from dicts.

PiperOrigin-RevId: 751446115

Author: manuadg-verily

jax_tpu_embedding/sparsecore/lib/core/input_preprocessing_cc_test.py

@@ -763,10 +763,7 @@ def test_multi_process_fdo(self, has_leading_dimension):
             allow_id_dropping=False,
         )
     )
-    stats = embedding.SparseDenseMatmulInputStats(
-        max_ids_per_partition=stats["max_ids"],
-        max_unique_ids_per_partition=stats["max_unique_ids"],
-    )
+    stats = embedding.SparseDenseMatmulInputStats.from_dict(stats)
     fdo_client.record(stats)
     fdo_client.publish()
     # Duplicated ids on row 0 and 6 are combined.

jax_tpu_embedding/sparsecore/lib/nn/embedding.py

@@ -68,6 +68,15 @@ class SparseDenseMatmulInputStats:
   max_ids_per_partition: Mapping[str, np.ndarray]
   max_unique_ids_per_partition: Mapping[str, np.ndarray]
 
+  @classmethod
+  def from_dict(
+      cls, stats: Mapping[str, Mapping[str, np.ndarray]]
+  ) -> "SparseDenseMatmulInputStats":
+    return cls(
+        max_ids_per_partition=stats["max_ids"],
+        max_unique_ids_per_partition=stats["max_unique_ids"],
+    )
+
 
 # TODO: b/346873239 - Add more checks for the feature specs to ensure all the
 # fields are valid.
@@ -371,10 +380,7 @@ def preprocess_sparse_dense_matmul_input(
 
   return SparseDenseMatmulInput(
       *preprocessed_inputs
-  ), SparseDenseMatmulInputStats(
-      max_ids_per_partition=stats["max_ids"],
-      max_unique_ids_per_partition=stats["max_unique_ids"],
-  )
+  ), SparseDenseMatmulInputStats.from_dict(stats)
 
 
 def _get_activation_for_feature(