Fix functional dict inputs to support optional ones (#21030)

* Fix functional dict inputs to support optional ones

* Add unit test for optional dict inputs

* Fix unit test formatting

Author: neo-alex

keras/src/models/functional.py

@@ -347,15 +347,15 @@ def shape_with_no_batch_size(x):
                 x[0] = None
             return tuple(x)
 
-        def make_spec_for_tensor(x):
+        def make_spec_for_tensor(x, name=None):
             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,
+                name=x._keras_history[0].name if name is None else name,
                 optional=optional,
             )
 
@@ -367,13 +367,7 @@ def make_spec_for_tensor(x):
                 # Case where `_nested_inputs` is a plain dict of Inputs.
                 names = sorted(self._inputs_struct.keys())
                 return [
-                    InputSpec(
-                        shape=shape_with_no_batch_size(
-                            self._inputs_struct[name].shape
-                        ),
-                        allow_last_axis_squeeze=True,
-                        name=name,
-                    )
+                    make_spec_for_tensor(self._inputs_struct[name], name=name)
                     for name in names
                 ]
             return None  # Deeply nested dict: skip checks.

keras/src/models/functional_test.py

@@ -574,6 +574,26 @@ def compute_output_shape(self, x_shape):
         self.assertAllClose(out, np.ones((2, 2)))
         # Note: it's not intended to work in symbolic mode (yet).
 
+    def test_optional_dict_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,), name="input1")
+        i2 = Input((2,), name="input2", optional=True)
+        outputs = OptionalInputLayer()(i1, i2)
+        model = Model({"input1": i1, "input2": i2}, outputs)
+
+        # Eager test
+        out = model({"input1": np.ones((2, 2)), "input2": None})
+        self.assertAllClose(out, np.ones((2, 2)))
+        # Note: it's not intended to work in symbolic mode (yet).
+
     def test_warning_for_mismatched_inputs_structure(self):
         def is_input_warning(w):
             return str(w.message).startswith(