Allow None inputs in Layer.build. (#21866)

If `call` takes a structure as inputs and some of the inputs are `None`, currently, building on call will fail.

Author: hertschuh

keras/src/layers/layer.py

@@ -1893,6 +1893,10 @@ def get_shapes_dict(call_spec):
     {"input_a_shape": (2, 3)}
     ```
     """
+
+    def standardize_shape_or_none(x):
+        return None if x is None else backend.standardize_shape(x.shape)
+
     shapes_dict = {}
     for k, v in call_spec.tensor_arguments_dict.items():
         if k == "mask" or k.endswith("_mask"):
@@ -1903,10 +1907,10 @@ def get_shapes_dict(call_spec):
             continue
         if k in call_spec.nested_tensor_argument_names:
             shapes_dict[f"{k}_shape"] = tree.map_structure(
-                lambda x: backend.standardize_shape(x.shape), v
+                standardize_shape_or_none, v
             )
         else:
-            shapes_dict[f"{k}_shape"] = backend.standardize_shape(v.shape)
+            shapes_dict[f"{k}_shape"] = standardize_shape_or_none(v)
     return shapes_dict
 
 

keras/src/layers/layer_test.py

@@ -516,6 +516,49 @@ def call(self, x1, x2):
         layer(x1=backend.KerasTensor((3, 4)), x2=backend.KerasTensor((3, 4)))
         self.assertLen(layer.weights, 4)
 
+        class DictLayerWithUnbuiltState(layers.Layer):
+            def __init__(self, units):
+                super().__init__()
+                self.dense = layers.Dense(units)
+
+            def call(self, xs):
+                result = self.dense(xs["x1"])
+                if xs.get("x2", None) is not None:
+                    result += self.dense(xs["x2"])
+                return result
+
+        layer = DictLayerWithUnbuiltState(2)
+        layer(
+            {
+                "x1": backend.KerasTensor((3, 4)),
+                "x2": backend.KerasTensor((3, 4)),
+            }
+        )
+        self.assertLen(layer.weights, 2)
+
+        layer = DictLayerWithUnbuiltState(2)
+        layer({"x1": backend.KerasTensor((3, 4)), "x2": None})
+        self.assertLen(layer.weights, 2)
+
+        class ListLayerWithUnbuiltState(layers.Layer):
+            def __init__(self, units):
+                super().__init__()
+                self.dense = layers.Dense(units)
+
+            def call(self, xs):
+                result = self.dense(xs[0])
+                if xs[1] is not None:
+                    result += self.dense(xs[1])
+                return result
+
+        layer = ListLayerWithUnbuiltState(2)
+        layer([backend.KerasTensor((3, 4)), backend.KerasTensor((3, 4))])
+        self.assertLen(layer.weights, 2)
+
+        layer = ListLayerWithUnbuiltState(2)
+        layer([backend.KerasTensor((3, 4)), None])
+        self.assertLen(layer.weights, 2)
+
     def test_activity_regularization(self):
         class ActivityRegularizer(layers.Layer):
             def call(self, x):