Fix support for optional inputs in model.fit

keras/src/backend/tensorflow/trainer.py

@@ -51,6 +51,7 @@ def distribute_reduction_method(self, value):
 
     def train_step(self, data):
         x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)
+        x = self._convert_optional_to_none(x)
 
         # Forward pass
         with tf.GradientTape() as tape:
@@ -86,6 +87,7 @@ def train_step(self, data):
 
     def test_step(self, data):
         x, y, sample_weight = data_adapter_utils.unpack_x_y_sample_weight(data)
+        x = self._convert_optional_to_none(x)
         if self._call_has_training_arg:
             y_pred = self(x, training=False)
         else:
@@ -101,12 +103,19 @@ def test_step(self, data):
 
     def predict_step(self, data):
         x, _, _ = data_adapter_utils.unpack_x_y_sample_weight(data)
+        x = self._convert_optional_to_none(x)
         if self._call_has_training_arg:
             y_pred = self(x, training=False)
         else:
             y_pred = self(x)
         return y_pred
 
+    def _convert_optional_to_none(self, x):
+        # Convert TF Optional implementations to None
+        return tree.map_structure(
+            lambda i: None if isinstance(i, tf.experimental.Optional) else i, x
+        )
+
     def _make_function(self, step_function):
         @tf.autograph.experimental.do_not_convert
         def one_step_on_data(data):

keras/src/trainers/data_adapters/array_data_adapter.py

@@ -76,7 +76,9 @@ def __init__(
         inputs = data_adapter_utils.pack_x_y_sample_weight(x, y, sample_weight)
 
         data_adapter_utils.check_data_cardinality(inputs)
-        num_samples = set(i.shape[0] for i in tree.flatten(inputs)).pop()
+        num_samples = set(
+            i.shape[0] for i in tree.flatten(inputs) if i is not None
+        ).pop()
         self._num_samples = num_samples
         self._inputs = inputs
 
@@ -269,7 +271,9 @@ def slice_and_convert(sliceable):
                     x = convert_to_tensor(x)
                     return x
 
-                return tree.map_structure(slice_and_convert, self.array)
+                return tree.map_structure(
+                    slice_and_convert, self.array, none_is_leaf=False
+                )
 
             def __len__(self):
                 return len(self.array[0])
@@ -337,7 +341,9 @@ def _get_iterator(self, slice_and_convert_fn, inputs):
             slice_indices_and_convert_fn = functools.partial(
                 slice_and_convert_fn, indices=indices
             )
-            yield tree.map_structure(slice_indices_and_convert_fn, inputs)
+            yield tree.map_structure(
+                slice_indices_and_convert_fn, inputs, none_is_leaf=False
+            )
 
     @property
     def num_batches(self):

keras/src/trainers/data_adapters/data_adapter_utils.py

@@ -101,7 +101,9 @@ def list_to_tuple(maybe_list):
 
 
 def check_data_cardinality(data):
-    num_samples = set(int(i.shape[0]) for i in tree.flatten(data))
+    num_samples = set(
+        int(i.shape[0]) for i in tree.flatten(data) if i is not None
+    )
     if len(num_samples) > 1:
         msg = (
             "Data cardinality is ambiguous. "
@@ -186,7 +188,9 @@ def get_single_tensor_spec(*tensors):
         else:
             return backend.KerasTensor(shape=shape, dtype=dtype)
 
-    return tree.map_structure(get_single_tensor_spec, *batches)
+    return tree.map_structure(
+        get_single_tensor_spec, *batches, none_is_leaf=False
+    )
 
 
 def convert_to_tf_tensor_spec(keras_tensor, batch_axis_to_none=True):
@@ -199,6 +203,8 @@ def convert_to_tf_tensor_spec(keras_tensor, batch_axis_to_none=True):
     """
     from keras.src.utils.module_utils import tensorflow as tf
 
+    if keras_tensor is None:
+        return tf.OptionalSpec(None)
     if not isinstance(keras_tensor, backend.KerasTensor):
         raise TypeError(
             f"Expected a KerasTensor, but got {keras_tensor} of type "
@@ -252,7 +258,9 @@ def convert_to_jax_compatible(x):
             return np.asarray(x)
 
     for batch in iterable:
-        yield tree.map_structure(convert_to_jax_compatible, batch)
+        yield tree.map_structure(
+            convert_to_jax_compatible, batch, none_is_leaf=False
+        )
 
 
 def get_numpy_iterator(iterable):
@@ -268,7 +276,7 @@ def convert_to_numpy(x):
         return x
 
     for batch in iterable:
-        yield tree.map_structure(convert_to_numpy, batch)
+        yield tree.map_structure(convert_to_numpy, batch, none_is_leaf=False)
 
 
 def get_torch_dataloader(iterable):
@@ -282,7 +290,9 @@ def __init__(self, iterable):
 
         def __iter__(self):
             for batch in self.iterable:
-                yield tree.map_structure(convert_to_tensor, batch)
+                yield tree.map_structure(
+                    convert_to_tensor, batch, none_is_leaf=False
+                )
 
     dataset = ConverterIterableDataset(iterable)
     # `batch_size=None` indicates that we should not re-batch

keras/src/trainers/data_adapters/generator_data_adapter.py

@@ -32,6 +32,8 @@ def get_tf_dataset(self):
         from keras.src.utils.module_utils import tensorflow as tf
 
         def convert_to_tf(x, spec):
+            if isinstance(spec, tf.OptionalSpec):
+                return x
             if data_adapter_utils.is_scipy_sparse(x):
                 x = data_adapter_utils.scipy_sparse_to_tf_sparse(x)
             elif data_adapter_utils.is_jax_sparse(x):
@@ -50,6 +52,14 @@ def convert_to_tf(x, spec):
 
         def get_tf_iterator():
             for batch in self.generator():
+                batch = tree.map_structure(
+                    (
+                        lambda i: tf.experimental.Optional.empty(None)
+                        if i is None
+                        else i
+                    ),
+                    batch,
+                )
                 batch = tree.map_structure(
                     convert_to_tf, batch, self._output_signature
                 )

keras/src/trainers/data_adapters/grain_dataset_adapter.py

@@ -80,7 +80,9 @@ def convert_to_numpy(x):
 
         class ConvertToNumpy(grain.transforms.Map):
             def map(self, x):
-                return tree.map_structure(convert_to_numpy, x)
+                return tree.map_structure(
+                    convert_to_numpy, x, none_is_leaf=False
+                )
 
         if isinstance(self._dataset, (grain.MapDataset, grain.IterDataset)):
             dataset = self._dataset.map(ConvertToNumpy())
@@ -109,7 +111,9 @@ def convert_to_jax_compatible(x):
 
         class ConvertToJaxCompatible(grain.transforms.Map):
             def map(self, x):
-                return tree.map_structure(convert_to_jax_compatible, x)
+                return tree.map_structure(
+                    convert_to_jax_compatible, x, none_is_leaf=False
+                )
 
         if isinstance(self._dataset, (grain.MapDataset, grain.IterDataset)):
             dataset = self._dataset.map(ConvertToJaxCompatible())
@@ -139,7 +143,7 @@ def convert_to_tf(x):
 
         class ConvertToTF(grain.transforms.Map):
             def map(self, x):
-                return tree.map_structure(convert_to_tf, x)
+                return tree.map_structure(convert_to_tf, x, none_is_leaf=False)
 
         # `tf.data.Dataset.from_generator` does not support lists as output.
         # We convert lists to tuples.

keras/src/trainers/data_adapters/tf_dataset_adapter.py

@@ -38,7 +38,9 @@ def get_numpy_iterator(self):
         from keras.src.backend.tensorflow.core import convert_to_numpy
 
         for batch in self._dataset:
-            yield tree.map_structure(convert_to_numpy, batch)
+            yield tree.map_structure(
+                convert_to_numpy, batch, none_is_leaf=False
+            )
 
     def get_jax_iterator(self):
         from keras.src.backend.tensorflow.core import convert_to_numpy
@@ -52,7 +54,7 @@ def convert_to_jax(x):
                 return convert_to_numpy(x)
 
         for batch in self._dataset:
-            yield tree.map_structure(convert_to_jax, batch)
+            yield tree.map_structure(convert_to_jax, batch, none_is_leaf=False)
 
     def get_tf_dataset(self):
         return self._dataset

keras/src/trainers/data_adapters/torch_data_loader_adapter.py

@@ -35,7 +35,9 @@ def get_numpy_iterator(self):
         for batch in self._dataloader:
             # shared memory using `np.asarray`
             yield tuple(
-                tree.map_structure(lambda x: np.asarray(x.cpu()), batch)
+                tree.map_structure(
+                    lambda x: np.asarray(x.cpu()), batch, none_is_leaf=False
+                )
             )
 
     def get_jax_iterator(self):

keras/src/tree/dmtree_impl.py

@@ -194,16 +194,32 @@ def flatten_with_path(structure):
     return flattened
 
 
-def map_structure(func, *structures):
+def map_structure(func, *structures, none_is_leaf=True):
     if not callable(func):
         raise TypeError(
             f"`func` must be callable, got {func} of type {type(func)}"
         )
 
+    map_func = func
+    if not none_is_leaf:
+
+        def func_skipping_none(*args):
+            # Check if the reference entry (first one) is None
+            if args[0] is None:
+                if not all(s is None for s in args):
+                    raise ValueError(
+                        "Structure mismatch: some arguments are None, others "
+                        "are not."
+                    )
+                return None
+            return func(*args)
+
+        map_func = func_skipping_none
+
     def func_traverse_wrapper(s):
         if is_nested(s):
             return None
-        ret = func(s)
+        ret = map_func(s)
         if ret is None:
             return dmtree.MAP_TO_NONE
         return ret
@@ -212,7 +228,7 @@ def func_traverse_wrapper(s):
         return traverse(func_traverse_wrapper, structures[0])
 
     with TypeErrorRemapping():
-        return dmtree.map_structure(func, *structures)
+        return dmtree.map_structure(map_func, *structures)
 
 
 def map_structure_up_to(shallow_structure, func, *structures):

keras/src/tree/optree_impl.py

@@ -93,14 +93,14 @@ def flatten_with_path(structure):
     return list(zip(paths, leaves))
 
 
-def map_structure(func, *structures):
+def map_structure(func, *structures, none_is_leaf=True):
     if not structures:
         raise ValueError("Must provide at least one structure")
 
     # Add check for same structures, otherwise optree just maps to shallowest.
     def func_with_check(*args):
         if not all(
-            optree.tree_is_leaf(s, none_is_leaf=True, namespace="keras")
+            optree.tree_is_leaf(s, none_is_leaf=none_is_leaf, namespace="keras")
             for s in args
         ):
             raise ValueError("Structures don't have the same nested structure.")
@@ -109,7 +109,7 @@ def func_with_check(*args):
     map_func = func_with_check if len(structures) > 1 else func
 
     return optree.tree_map(
-        map_func, *structures, none_is_leaf=True, namespace="keras"
+        map_func, *structures, none_is_leaf=none_is_leaf, namespace="keras"
     )
 
 

keras/src/tree/tree_api.py

@@ -160,7 +160,7 @@ def flatten_with_path(structure):
 
 
 @keras_export("keras.tree.map_structure")
-def map_structure(func, *structures):
+def map_structure(func, *structures, none_is_leaf=True):
     """Maps `func` through given structures.
 
     Examples:
@@ -179,6 +179,7 @@ def map_structure(func, *structures):
     Args:
         func: A callable that accepts as many arguments as there are structures.
         *structures: Arbitrarily nested structures of the same layout.
+        none_is_leaf: If True, None is treated as a leaf.
 
     Returns:
         A new structure with the same layout as the given ones.
@@ -189,7 +190,7 @@ def map_structure(func, *structures):
             the nested structures don't match according to the rules of
             `assert_same_structure`.
     """
-    return tree_impl.map_structure(func, *structures)
+    return tree_impl.map_structure(func, *structures, none_is_leaf=none_is_leaf)
 
 
 @keras_export("keras.tree.map_structure_up_to")