Add support for optional Functional inputs.

Author: fchollet

keras/src/layers/core/input_layer.py

@@ -16,6 +16,7 @@ def __init__(
         sparse=None,
         batch_shape=None,
         input_tensor=None,
+        optional=False,
         name=None,
         **kwargs,
     ):
@@ -69,6 +70,7 @@ def __init__(
         self._input_tensor = input_tensor
         Node(operation=self, call_args=(), call_kwargs={}, outputs=input_tensor)
         self.built = True
+        self.optional = optional
 
     def call(self):
         return
@@ -95,6 +97,7 @@ def Input(
     batch_shape=None,
     name=None,
     tensor=None,
+    optional=False,
 ):
     """Used to instantiate a Keras tensor.
 
@@ -127,6 +130,8 @@ def Input(
         tensor: Optional existing tensor to wrap into the `Input` layer.
             If set, the layer will use this tensor rather
             than creating a new placeholder tensor.
+        optional: Boolean, whether the input is optional or not.
+            An optional input can accept `None` values.
 
     Returns:
       A Keras tensor.
@@ -148,5 +153,6 @@ def Input(
         batch_shape=batch_shape,
         name=name,
         input_tensor=tensor,
+        optional=optional,
     )
     return layer.output

keras/src/layers/input_spec.py

@@ -30,6 +30,8 @@ class InputSpec:
             as long as the last axis of the spec is 1.
         name: Expected key corresponding to this input when passing data as
             a dictionary.
+        optional: Boolean, whether the input is optional or not.
+            An optional input can accept `None` values.
 
     Example:
 
@@ -56,6 +58,7 @@ def __init__(
         axes=None,
         allow_last_axis_squeeze=False,
         name=None,
+        optional=False,
     ):
         self.dtype = (
             backend.standardize_dtype(dtype) if dtype is not None else None
@@ -69,6 +72,7 @@ def __init__(
         self.max_ndim = max_ndim
         self.min_ndim = min_ndim
         self.name = name
+        self.optional = optional
         self.allow_last_axis_squeeze = allow_last_axis_squeeze
         try:
             axes = axes or {}
@@ -152,12 +156,18 @@ def assert_input_compatibility(input_spec, inputs, layer_name):
             inputs = list_inputs
 
     inputs = tree.flatten(inputs)
-    if len(input_spec) != len(inputs):
+    if len(inputs) != len(input_spec):
         raise ValueError(
-            f"Layer '{layer_name}' expected {len(input_spec)} input(s). "
-            f"Received {len(inputs)} instead."
+            f'Layer "{layer_name}" expects {len(input_spec)} input(s),'
+            f" but it received {len(inputs)} input tensors. "
+            f"Inputs received: {inputs}"
         )
-    for x in inputs:
+    for input_index, (x, spec) in enumerate(zip(inputs, input_spec)):
+        if spec is None:
+            continue
+        if x is None and spec.optional:
+            continue
+
         # Having a shape/dtype is the only commonality of the various
         # tensor-like objects that may be passed. The most common kind of
         # invalid type we are guarding for is a Layer instance (Functional API),
@@ -168,16 +178,6 @@ def assert_input_compatibility(input_spec, inputs, layer_name):
                 f"(of type {type(x)}) as input for layer '{layer_name}'."
             )
 
-    if len(inputs) != len(input_spec):
-        raise ValueError(
-            f'Layer "{layer_name}" expects {len(input_spec)} input(s),'
-            f" but it received {len(inputs)} input tensors. "
-            f"Inputs received: {inputs}"
-        )
-    for input_index, (x, spec) in enumerate(zip(inputs, input_spec)):
-        if spec is None:
-            continue
-
         shape = backend.standardize_shape(x.shape)
         ndim = len(shape)
         # Check ndim.

keras/src/layers/layer.py

@@ -747,8 +747,10 @@ def maybe_convert(x):
         # 2. Enforce that only tensors can be passed positionally.
         if not self._allow_non_tensor_positional_args:
             for arg in tree.flatten(args):
-                if not isinstance(arg, KerasTensor) and not backend.is_tensor(
-                    arg
+                if (
+                    not isinstance(arg, KerasTensor)
+                    and not backend.is_tensor(arg)
+                    and arg is not None
                 ):
                     raise ValueError(
                         "Only input tensors may be passed as "

keras/src/models/functional.py

@@ -207,15 +207,23 @@ def _flatten_to_reference_inputs(self, inputs):
     def _convert_inputs_to_tensors(self, flat_inputs):
         converted = []
         for x, input in zip(flat_inputs, self._inputs):
-            converted.append(
-                ops.convert_to_tensor(x, dtype=input.dtype, sparse=input.sparse)
-            )
+            if x is None:  # TODO: check if optional
+                converted.append(x)
+            else:
+                converted.append(
+                    ops.convert_to_tensor(
+                        x, dtype=input.dtype, sparse=input.sparse
+                    )
+                )
         return converted
 
     def _adjust_input_rank(self, flat_inputs):
         flat_ref_shapes = [x.shape for x in self._inputs]
         adjusted = []
         for x, ref_shape in zip(flat_inputs, flat_ref_shapes):
+            if x is None:
+                adjusted.append(x)
+                continue
             x_rank = len(x.shape)
             ref_rank = len(ref_shape)
             if x_rank == ref_rank:
@@ -273,10 +281,15 @@ def shape_with_no_batch_size(x):
             return tuple(x)
 
         def make_spec_for_tensor(x):
+            optional = False
+            if isinstance(x._keras_history[0], InputLayer):
+                if x._keras_history[0].optional:
+                    optional = True
             return InputSpec(
                 shape=shape_with_no_batch_size(x.shape),
                 allow_last_axis_squeeze=True,
                 name=x._keras_history[0].name,
+                optional=optional,
             )
 
         if isinstance(self._inputs_struct, dict):

keras/src/models/functional_test.py

@@ -304,7 +304,7 @@ def test_bad_input_spec(self):
             ValueError, r"expected shape=\(None, 4\), found shape=\(2, 3\)"
         ):
             model(np.zeros((2, 3)))
-        with self.assertRaisesRegex(ValueError, "expected 1 input"):
+        with self.assertRaisesRegex(ValueError, "expects 1 input"):
             model([np.zeros((2, 4)), np.zeros((2, 4))])
 
         # List input
@@ -313,7 +313,7 @@ def test_bad_input_spec(self):
         x = input_a + input_b
         outputs = layers.Dense(2)(x)
         model = Functional([input_a, input_b], outputs)
-        with self.assertRaisesRegex(ValueError, "expected 2 input"):
+        with self.assertRaisesRegex(ValueError, "expects 2 input"):
             model(np.zeros((2, 3)))
         with self.assertRaisesRegex(
             ValueError, r"expected shape=\(None, 4\), found shape=\(2, 3\)"
@@ -322,7 +322,7 @@ def test_bad_input_spec(self):
 
         # Dict input
         model = Functional({"a": input_a, "b": input_b}, outputs)
-        with self.assertRaisesRegex(ValueError, "expected 2 input"):
+        with self.assertRaisesRegex(ValueError, "expects 2 input"):
             model(np.zeros((2, 3)))
         with self.assertRaisesRegex(
             ValueError, r"expected shape=\(None, 4\), found shape=\(2, 3\)"
@@ -432,6 +432,26 @@ def test_deeply_nested_model(self):
             out_eager["others"]["3"], new_out_eager["others"]["3"]
         )
 
+    def test_optional_inputs(self):
+        class OptionalInputLayer(layers.Layer):
+            def call(self, x, y=None):
+                if y is not None:
+                    return x + y
+                return x
+
+            def compute_output_shape(self, x_shape):
+                return x_shape
+
+        i1 = Input((2,))
+        i2 = Input((2,), optional=True)
+        outputs = OptionalInputLayer()(i1, i2)
+        model = Model([i1, i2], outputs)
+
+        # Eager test
+        out = model([np.ones((2, 2)), None])
+        self.assertAllClose(out, np.ones((2, 2)))
+        # Note: it's not intended to work in symbolic mode (yet).
+
     def test_add_loss(self):
         # TODO
         pass

keras/src/ops/operation.py

@@ -79,19 +79,16 @@ def compute_output_spec(self, *args, **kwargs):
         try:
             return backend.compute_output_spec(self.call, *args, **kwargs)
         except Exception as e:
-            if isinstance(e, TypeError):
-                raise e
-            else:
-                new_e = RuntimeError(
-                    "Could not automatically infer the output shape / dtype of "
-                    f"'{self.name}' (of type {self.__class__.__name__}). "
-                    f"Either the `{self.__class__.__name__}.call()` method "
-                    f"is incorrect, or you need to implement the "
-                    f"`{self.__class__.__name__}.compute_output_spec() / "
-                    "compute_output_shape()` method. "
-                    f"Error encountered:\n\n{e}"
-                )
-                raise new_e.with_traceback(e.__traceback__) from None
+            new_e = e.__class__(
+                "Could not automatically infer the output shape / dtype of "
+                f"'{self.name}' (of type {self.__class__.__name__}). "
+                f"Either the `{self.__class__.__name__}.call()` method "
+                f"is incorrect, or you need to implement the "
+                f"`{self.__class__.__name__}.compute_output_spec() / "
+                "compute_output_shape()` method. "
+                f"Error encountered:\n\n{e}"
+            )
+            raise new_e.with_traceback(e.__traceback__) from None
 
     def __new__(cls, *args, **kwargs):
         """We override __new__ to saving serializable constructor arguments.

keras/src/ops/symbolic_arguments.py

@@ -40,9 +40,7 @@ def fill_in(self, tensor_dict):
 
         def switch_fn(x):
             if isinstance(x, KerasTensor):
-                val = tensor_dict.get(id(x), None)
-                if val is not None:
-                    return val
+                return tensor_dict.get(id(x), None)
             return x
 
         return self.convert(switch_fn)

keras/src/ops/symbolic_arguments_test.py

@@ -99,8 +99,7 @@ def test_fill_in_multiple_arg(self):
 
         # Call the method to be tested
         result, _ = sym_args.fill_in(tensor_dict)
-
-        self.assertEqual(result, ((a, 2),))
+        self.assertEqual(result, ((None, 2),))
 
     # Testing fill in function for args and kwargs
     def test_fill_in(self):
@@ -115,9 +114,8 @@ def test_fill_in(self):
                 a,
                 b,
             ),
-            {1: c},
+            {"1": c},
         )
 
         (values, _) = sym_args.fill_in(dictionary)
-
-        self.assertEqual(values, ((3, b), {1: 2}))
+        self.assertEqual(values, ((3, None), {"1": 2}))