Add batch_number to AbstractInputBatch.

This CL adds the batch_number attribute to the AbstractInputBatch base class. The number should be unique and incremental, but can be reset to 0 on restart or between epochs.

PiperOrigin-RevId: 782134551

Author: patnotz

recml/layers/linen/sparsecore.py

@@ -328,6 +328,8 @@ def _to_np(x: Any) -> np.ndarray:
         if weights[key] is not None:
           weights[key] = np.reshape(weights[key], (-1, 1))
 
+    # TODO(patn): Find the step number. DO NOT SUBMIT
+    batch_number = 0
     csr_inputs, _ = embedding.preprocess_sparse_dense_matmul_input(
         features=features,
         features_weights=weights,
@@ -337,6 +339,7 @@ def _to_np(x: Any) -> np.ndarray:
         num_sc_per_device=self.sparsecore_config.num_sc_per_device,
         sharding_strategy=self.sparsecore_config.sharding_strategy,
         allow_id_dropping=False,
+        batch_number=batch_number,
     )
 
     processed_inputs = {