Use the element spec to get the input shapes instead of getting a batch.

PiperOrigin-RevId: 764451548

Author: Hilly12

recml/core/training/core.py

@@ -174,9 +174,9 @@ def get_iterators(
 
 
 def get_shape(
-    x: tf.Tensor | tf.SparseTensor | tf.RaggedTensor,
+    x: tf.Tensor | tf.SparseTensor | tf.RaggedTensor | tf.TensorSpec,
 ) -> Sequence[int | None]:
-  """Gets the shape of a dense / sparse / ragged tensor."""
+  """Gets the shape of a dense / sparse / ragged tensor or tensor spec."""
   if isinstance(x, tf.SparseTensor):
     return [x.shape[0]] + [None for _ in x.shape[1:]]
   return x.shape.as_list()

recml/core/training/keras_trainer.py

@@ -106,7 +106,9 @@ def __init__(
       self,
       *,
       distribution: (
-          keras.distribution.DataParallel | keras.distribution.ModelParallel
+          keras.distribution.DataParallel
+          | keras.distribution.ModelParallel
+          | None
       ) = None,
       model_dir: str | None = None,
       train_steps: int = 0,
@@ -128,10 +130,7 @@ def __init__(
     # This should be set before any layers are constructed and this is a
     # fallback in case the trainer binary doesn't already do this.
     if (
-        isinstance(
-            distribution,
-            (keras.distribution.DataParallel, keras.distribution.ModelParallel),
-        )
+        distribution is not None
         and keras.distribution.distribution() != distribution
     ):
       if hasattr(distribution, "_auto_shard_dataset"):
@@ -175,6 +174,7 @@ def __init__(
           ),
       ]
     else:
+      self._checkpoint_manager = None
       self._train_callbacks = [
           keras.callbacks.TensorBoard(
               log_dir=os.path.join(model_dir, core.LOG_DIR),
@@ -199,13 +199,13 @@ def __init__(
       ]
 
   def _maybe_get_model_kws(
-      self, task: KerasTask, dataset: keras.Model
+      self, task: KerasTask, dataset: tf.data.Dataset
   ) -> Mapping[str, Any]:
     kws = {}
     if py_utils.has_argument(task.create_model, "input_shapes"):
-      batch = next(iter(dataset))
-      x, *_ = keras.utils.unpack_x_y_sample_weight(batch)
-      kws["input_shapes"]: keras.tree.map_structure(core.get_shape, x)  # pylint: disable=undefined-variable
+      batch_spec = dataset.element_spec
+      x, *_ = keras.utils.unpack_x_y_sample_weight(batch_spec)
+      kws["input_shapes"] = keras.tree.map_structure(core.get_shape, x)
 
     return kws