tf-keras-nightly 2.21.0.dev2025102309__py3-none-any.whl → 2.21.0.dev2025120310__py3-none-any.whl

tf_keras/__init__.py

@@ -27,4 +27,4 @@ from tf_keras.src.engine.sequential import Sequential
 from tf_keras.src.engine.training import Model
 
 
-__version__ = "2.21.0.dev2025102309"
+__version__ = "2.21.0.dev2025120310"

tf_keras/src/engine/data_adapter.py

@@ -231,7 +231,7 @@ class TensorLikeDataAdapter(DataAdapter):
                 return True
             return False
 
-        return all(_is_tensor(v) for v in flat_inputs)
+        return all(_is_tensor(v) for v in flat_inputs if v is not None)
 
     def __init__(
         self,
@@ -259,7 +259,7 @@ class TensorLikeDataAdapter(DataAdapter):
         inputs = pack_x_y_sample_weight(x, y, sample_weights)
 
         num_samples = set(
-            int(i.shape[0]) for i in tf.nest.flatten(inputs)
+            int(i.shape[0]) for i in tf.nest.flatten(inputs) if i is not None
         ).pop()
         _check_data_cardinality(inputs)
 
@@ -386,7 +386,7 @@ class TensorLikeDataAdapter(DataAdapter):
 
         def grab_batch(i, data):
             return tf.nest.map_structure(
-                lambda d: tf.gather(d, i, axis=0), data
+                lambda d: tf.gather(d, i, axis=0) if d is not None else d, data
             )
 
         dataset = dataset.map(grab_batch, num_parallel_calls=tf.data.AUTOTUNE)
@@ -459,7 +459,7 @@ class GenericArrayLikeDataAdapter(TensorLikeDataAdapter):
         if not TensorLikeDataAdapter.can_handle(
             x, y
         ) and not CompositeTensorDataAdapter.can_handle(x, y):
-            return all(_is_array_like(v) for v in flat_inputs)
+            return all(_is_array_like(v) for v in flat_inputs if v is not None)
         else:
             return False
 
@@ -496,7 +496,7 @@ class GenericArrayLikeDataAdapter(TensorLikeDataAdapter):
             shape[0] = None
             return tuple(shape)
 
-        flat_dtypes = [inp.dtype for inp in flat_inputs]
+        flat_dtypes = [inp.dtype for inp in flat_inputs if inp is not None]
         contiguous = True
         if self._shuffle and self._shuffle != "batch":
             contiguous = False
@@ -509,15 +509,26 @@ class GenericArrayLikeDataAdapter(TensorLikeDataAdapter):
             # to a Tensor may force it into memory..
             def py_method(ind):
                 def slice_array(data):
+                    if data is None:
+                        return None
                     return training_utils.slice_arrays(
                         data, ind.numpy(), contiguous=contiguous
                     )
 
-                return [slice_array(inp) for inp in flat_inputs]
+                return [
+                    slice_array(inp) for inp in flat_inputs if inp is not None
+                ]
 
-            flat_out = tf.py_function(py_method, [indices], flat_dtypes)
-            for v, original_inp in zip(flat_out, flat_inputs):
-                v.set_shape(dynamic_shape_like(original_inp))
+            results = tf.py_function(py_method, [indices], flat_dtypes)
+            results_it = iter(results)
+            flat_out = []
+            for original_inp in flat_inputs:
+                if original_inp is None:
+                    flat_out.append(None)
+                else:
+                    v = next(results_it)
+                    v.set_shape(dynamic_shape_like(original_inp))
+                    flat_out.append(v)
             return tf.nest.pack_sequence_as(inputs, flat_out)
 
         dataset = indices_dataset.map(
@@ -608,8 +619,10 @@ class CompositeTensorDataAdapter(DataAdapter):
                 return True
             return _is_composite(v)
 
-        return any(_is_composite(v) for v in flat_inputs) and all(
-            _is_tensor_or_composite(v) for v in flat_inputs
+        return any(
+            _is_composite(v) for v in flat_inputs if v is not None
+        ) and all(
+            _is_tensor_or_composite(v) for v in flat_inputs if v is not None
         )
 
     def __init__(
@@ -1944,14 +1957,18 @@ def single_batch_iterator(
 
 
 def _check_data_cardinality(data):
-    num_samples = set(int(i.shape[0]) for i in tf.nest.flatten(data))
+    num_samples = set(
+        int(i.shape[0]) for i in tf.nest.flatten(data) if i is not None
+    )
     if len(num_samples) > 1:
         msg = "Data cardinality is ambiguous:\n"
         for label, single_data in zip(["x", "y", "sample_weight"], data):
             msg += "  {} sizes: {}\n".format(
                 label,
                 ", ".join(
-                    str(i.shape[0]) for i in tf.nest.flatten(single_data)
+                    str(i.shape[0])
+                    for i in tf.nest.flatten(single_data)
+                    if i is not None
                 ),
             )
         msg += "Make sure all arrays contain the same number of samples."

tf_keras/src/engine/functional.py

@@ -351,25 +351,45 @@ class Functional(training_lib.Model):
         if isinstance(self._nested_inputs, dict):
             # Case where `_nested_inputs` is a plain dict of Inputs.
             names = sorted(self._nested_inputs.keys())
-            return [
-                input_spec.InputSpec(
-                    shape=shape_with_no_batch_size(self._nested_inputs[name]),
-                    allow_last_axis_squeeze=True,
-                    name=name,
+            specs = []
+            for name in names:
+                layer = self._nested_inputs[name]._keras_history.layer
+                optional = (
+                    layer.optional
+                    if isinstance(layer, input_layer_module.InputLayer)
+                    else False
                 )
-                for name in names
-            ]
+                specs.append(
+                    input_spec.InputSpec(
+                        shape=shape_with_no_batch_size(
+                            self._nested_inputs[name]
+                        ),
+                        allow_last_axis_squeeze=True,
+                        name=name,
+                        optional=optional,
+                    )
+                )
+            return specs
         else:
             # Single input, or list / tuple of inputs.
             # The data may be passed as a dict keyed by input name.
-            return [
-                input_spec.InputSpec(
-                    shape=shape_with_no_batch_size(x),
-                    allow_last_axis_squeeze=True,
-                    name=x._keras_history.layer.name,
+            specs = []
+            for x in self.inputs:
+                layer = x._keras_history.layer
+                optional = (
+                    layer.optional
+                    if isinstance(layer, input_layer_module.InputLayer)
+                    else False
                 )
-                for x in self.inputs
-            ]
+                specs.append(
+                    input_spec.InputSpec(
+                        shape=shape_with_no_batch_size(x),
+                        allow_last_axis_squeeze=True,
+                        name=x._keras_history.layer.name,
+                        optional=optional,
+                    )
+                )
+            return specs
 
     @input_spec.setter
     def input_spec(self, value):
@@ -644,7 +664,8 @@ class Functional(training_lib.Model):
         else:
             masks = self._flatten_to_reference_inputs(mask)
         for input_t, mask in zip(inputs, masks):
-            input_t._keras_mask = mask
+            if input_t is not None:
+                input_t._keras_mask = mask
 
         # Dictionary mapping reference tensors to computed tensors.
         tensor_dict = {}

tf_keras/src/engine/input_layer.py

@@ -98,6 +98,8 @@ class InputLayer(base_layer.Layer):
             `tf.TypeSpec` represents the entire batch. When provided, all other
             args except name must be `None`.
         name: Optional name of the layer (string).
+        optional: Boolean, whether the input is optional or not. An optional
+            input can accept `None` values.
     """
 
     @traceback_utils.filter_traceback
@@ -111,6 +113,7 @@ class InputLayer(base_layer.Layer):
         name=None,
         ragged=None,
         type_spec=None,
+        optional=False,
         **kwargs,
     ):
         self._init_input_shape = input_shape
@@ -180,6 +183,7 @@ class InputLayer(base_layer.Layer):
         self.ragged = True if ragged else False
         self.batch_size = batch_size
         self.supports_masking = True
+        self.optional = optional
 
         if isinstance(input_shape, tf.TensorShape):
             input_shape = tuple(input_shape.as_list())
@@ -284,6 +288,7 @@ class InputLayer(base_layer.Layer):
                 "sparse": self.sparse,
                 "ragged": self.ragged,
                 "name": self.name,
+                "optional": self.optional,
             }
         return config
 
@@ -303,6 +308,7 @@ def Input(
     tensor=None,
     ragged=None,
     type_spec=None,
+    optional=False,
     **kwargs,
 ):
     """`Input()` is used to instantiate a TF-Keras tensor.
@@ -341,6 +347,8 @@ def Input(
             [this guide](https://www.tensorflow.org/guide/ragged_tensor).
         type_spec: A `tf.TypeSpec` object to create the input placeholder from.
             When provided, all other args except name must be None.
+        optional: Boolean, whether the input is optional or not. An optional
+            input can accept `None` values.
         **kwargs: deprecated arguments support. Supports `batch_shape` and
             `batch_input_shape`.
 
@@ -415,6 +423,7 @@ def Input(
         "ragged": ragged,
         "input_tensor": tensor,
         "type_spec": type_spec,
+        "optional": optional,
     }
 
     batch_input_shape = kwargs.pop(

tf_keras/src/engine/input_spec.py

@@ -56,6 +56,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:
 
@@ -82,6 +84,7 @@ class InputSpec:
         axes=None,
         allow_last_axis_squeeze=False,
         name=None,
+        optional=False,
     ):
         self.dtype = tf.as_dtype(dtype).name if dtype is not None else None
         shape = tf.TensorShape(shape)
@@ -99,6 +102,7 @@ class InputSpec:
         self.min_ndim = min_ndim
         self.name = name
         self.allow_last_axis_squeeze = allow_last_axis_squeeze
+        self.optional = optional
         try:
             axes = axes or {}
             self.axes = {int(k): axes[k] for k in axes}
@@ -204,7 +208,11 @@ def assert_input_compatibility(input_spec, inputs, layer_name):
             inputs = list_inputs
 
     inputs = tf.nest.flatten(inputs)
-    for x in inputs:
+    for _, (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),
@@ -224,6 +232,8 @@ def assert_input_compatibility(input_spec, inputs, layer_name):
     for input_index, (x, spec) in enumerate(zip(inputs, input_spec)):
         if spec is None:
             continue
+        if x is None and spec.optional:
+            continue
 
         shape = tf.TensorShape(x.shape)
         if shape.rank is None:

tf_keras/src/layers/activation/softmax.py

@@ -70,6 +70,8 @@ class Softmax(Layer):
     Args:
         axis: Integer, or list of Integers, axis along which the softmax
             normalization is applied.
+        robust_masking: Bool, if true will use a more robust implementation when
+            dealing with masks.
     Call arguments:
         inputs: The inputs, or logits to the softmax layer.
         mask: A boolean mask of the same shape as `inputs`. The mask
@@ -80,23 +82,34 @@ class Softmax(Layer):
         Softmaxed output with the same shape as `inputs`.
     """
 
-    def __init__(self, axis=-1, **kwargs):
+    def __init__(self, axis=-1, robust_masking=False, **kwargs):
         super().__init__(**kwargs)
         self.supports_masking = True
+        self.robust_masking = robust_masking
         self.axis = axis
 
     def call(self, inputs, mask=None):
         if mask is not None:
-            # Since mask is 1.0 for positions we want to keep and 0.0 for masked
-            # positions, this operation will create a tensor which is 0.0 for
-            # positions we want to attend and -1e.9 for masked positions.
-            adder = (1.0 - tf.cast(mask, inputs.dtype)) * (
-                _large_compatible_negative(inputs.dtype)
-            )
-
-            # Since we are adding it to the raw scores before the softmax, this
-            # is effectively the same as removing these entirely.
-            inputs += adder
+            if self.robust_masking:
+                # We keep the positions where the mask is True or > 0.5, and set
+                # the other (masked) positions to -1e.9.
+                if mask.dtype is not tf.bool:
+                    mask = tf.greater(mask, tf.constant(0.5, dtype=mask.dtype))
+                inputs = tf.where(
+                    mask, inputs, _large_compatible_negative(inputs.dtype)
+                )
+            else:
+                # Since mask is 1.0 for positions we want to keep and 0.0 for
+                # masked positions, this operation will create a tensor which is
+                # 0.0 for positions we want to attend and -1e.9 for masked
+                # positions.
+                adder = (1.0 - tf.cast(mask, inputs.dtype)) * (
+                    _large_compatible_negative(inputs.dtype)
+                )
+
+                # Since we are adding it to the raw scores before the softmax,
+                # this is effectively the same as removing these entirely.
+                inputs += adder
         if isinstance(self.axis, (tuple, list)):
             if len(self.axis) > 1:
                 return tf.exp(
@@ -109,6 +122,8 @@ class Softmax(Layer):
 
     def get_config(self):
         config = {"axis": self.axis}
+        if self.robust_masking:
+            config["robust_masking"] = True
         base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
 

tf_keras/src/layers/attention/multi_head_attention.py

@@ -198,6 +198,8 @@ class MultiHeadAttention(Layer):
         activity_regularizer: Regularizer for dense layer activity.
         kernel_constraint: Constraint for dense layer kernels.
         bias_constraint: Constraint for dense layer kernels.
+        softmax_robust_masking: If true will use a more numerically robust
+            masking impl.
 
     Call arguments:
         query: Query `Tensor` of shape `(B, T, dim)`.
@@ -247,6 +249,7 @@ class MultiHeadAttention(Layer):
         activity_regularizer=None,
         kernel_constraint=None,
         bias_constraint=None,
+        softmax_robust_masking=False,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -264,6 +267,7 @@ class MultiHeadAttention(Layer):
         self._activity_regularizer = regularizers.get(activity_regularizer)
         self._kernel_constraint = constraints.get(kernel_constraint)
         self._bias_constraint = constraints.get(bias_constraint)
+        self._softmax_robust_masking = softmax_robust_masking
         if attention_axes is not None and not isinstance(
             attention_axes, collections.abc.Sized
         ):
@@ -298,6 +302,7 @@ class MultiHeadAttention(Layer):
             "query_shape": self._query_shape,
             "key_shape": self._key_shape,
             "value_shape": self._value_shape,
+            "softmax_robust_masking": self._softmax_robust_masking,
         }
         base_config = super().get_config()
         return dict(list(base_config.items()) + list(config.items()))
@@ -476,7 +481,9 @@ class MultiHeadAttention(Layer):
             )
         )
         self._softmax = activation.Softmax(
-            axis=norm_axes, dtype=self._dtype_policy
+            axis=norm_axes,
+            robust_masking=self._softmax_robust_masking,
+            dtype=self._dtype_policy,
         )
         self._dropout_layer = regularization.Dropout(
             rate=self._dropout, dtype=self._dtype_policy

{tf_keras_nightly-2.21.0.dev2025102309.dist-info → tf_keras_nightly-2.21.0.dev2025120310.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: tf_keras-nightly
-Version: 2.21.0.dev2025102309
+Version: 2.21.0.dev2025120310
 Summary: Deep learning for humans.
 Home-page: https://keras.io/
 Download-URL: https://github.com/keras-team/tf-keras/tags

{tf_keras_nightly-2.21.0.dev2025102309.dist-info → tf_keras_nightly-2.21.0.dev2025120310.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tf_keras/__init__.py,sha256=9UbilJSxb-fHpU9Mwg1YTRm402E8QjZfaZxDRCmExBg,911
+tf_keras/__init__.py,sha256=S8NBm8tTVyOyDqCg4ZaB1hQ5KwSr84xqdTQ-ni4f6Ws,911
 tf_keras/__internal__/__init__.py,sha256=OHQbeIC0QtRBI7dgXaJaVbH8F00x8dCI-DvEcIfyMsE,671
 tf_keras/__internal__/backend/__init__.py,sha256=LnMs2A6685gDG79fxqmdulIYlVE_3WmXlBTBo9ZWYcw,162
 tf_keras/__internal__/layers/__init__.py,sha256=F5SGMhOTPzm-PR44VrfinURHcVeQPIEdwnZlAkSTB3A,176
@@ -276,11 +276,11 @@ tf_keras/src/engine/base_layer_utils.py,sha256=AFjqwXM-WShf0dfsyIotlXYIRJlqYyjQh
 tf_keras/src/engine/base_layer_v1.py,sha256=MMfdUKB8tjbjjX9Pj5b6E5XgrM-BnVx0ilSndcR_3QA,102875
 tf_keras/src/engine/base_preprocessing_layer.py,sha256=xne5VVtj9_IE1_cjh-kaPk-utoMY7mYwTOcgybFfY34,12650
 tf_keras/src/engine/compile_utils.py,sha256=F6KxbaXnppns5XCOJl8wzsiQ1riEp43s0G0SWsWAUE0,31757
-tf_keras/src/engine/data_adapter.py,sha256=UqYJBUDiS-vyu7euVYxQrXw0U9-piO7SwTetkGBSMwg,71654
-tf_keras/src/engine/functional.py,sha256=igoSYLFSVYvZ8EUjwbowrDyYywI-V3h9ypeybIN_Jr4,70182
+tf_keras/src/engine/data_adapter.py,sha256=N5UV4KYF-F7YJdB5kHR8pACSrFlSqQte_DsmA8Ksa6Y,72257
+tf_keras/src/engine/functional.py,sha256=ojvj0DkGrnqd16dEIqj2AgFlmxB2s6e-3qjs78fN85E,71006
 tf_keras/src/engine/functional_utils.py,sha256=5creFfo9UoG5OLJgkcw9gsfT-qch-RamT5IsU8675rU,11048
-tf_keras/src/engine/input_layer.py,sha256=QVAA9ZrhfUlcx0Tj_UuNF3t1nxYrhyks6vDJJeb18W8,18258
-tf_keras/src/engine/input_spec.py,sha256=W3mojApaM_lN8Vr2MCvddE8RhHckxwObzpySPm16oEM,12076
+tf_keras/src/engine/input_layer.py,sha256=SRoRYG_PElMRMAypjeQRZ_0Ub4tu5jdRc6ASL59BAsY,18650
+tf_keras/src/engine/input_spec.py,sha256=H2U8yNz7eabUozm4QAcL3XcQoxj6iAwvK6ecU7w8O5g,12455
 tf_keras/src/engine/keras_tensor.py,sha256=rmIyf-sMKzGAMXzob0hCTZ3qA4JBYyIM85XUdmOPmqQ,28858
 tf_keras/src/engine/node.py,sha256=mevKNFEtzeVbwLRuwB7sMzQGKt6ppIxLmMcfQMzu8N8,14254
 tf_keras/src/engine/partial_batch_padding_handler.py,sha256=TNZvGXL-fvmZLLHIMPX_hy0w9LT8W52DHW7ZtnEvBvI,4325
@@ -314,13 +314,13 @@ tf_keras/src/layers/activation/elu.py,sha256=n-WAE6NjC9mbqcV7Kxgpt8tTbvwCQIGsoCV
 tf_keras/src/layers/activation/leaky_relu.py,sha256=cJmpwgg4KEu--iK9gFuJT7uEGpDArB8q-XNBmJfC7_U,2618
 tf_keras/src/layers/activation/prelu.py,sha256=D2yhneQrYQP6aHSK8nvnMKa1hIeuPZO_XCB2Cu9Cl4Y,4440
 tf_keras/src/layers/activation/relu.py,sha256=JklQuReRiR3huAGr3QRtuGL0URpdspDFzBNjZgv0HDw,4281
-tf_keras/src/layers/activation/softmax.py,sha256=G6MfTCogGTKwyP7b6ByxeIHFNQtUKgrZXB8MP9hNstQ,4105
+tf_keras/src/layers/activation/softmax.py,sha256=0g8uN5N8QDW8lj6nGabR-EBk58njbiNdhDzglv9rxXU,4861
 tf_keras/src/layers/activation/thresholded_relu.py,sha256=rQLn9cr-w6hVJET2mS7OIQ9diiUiqUrX4CysXKNYbmg,2503
 tf_keras/src/layers/attention/__init__.py,sha256=6HjPSyLhs_bf4erT65KyhSCHQF7WeWZe9YTH7iW6Nek,945
 tf_keras/src/layers/attention/additive_attention.py,sha256=jie0cAXJEjU4xXK_Ur1SrEL9RqDIIAPyaAkK8O71TEs,7485
 tf_keras/src/layers/attention/attention.py,sha256=TCnoOWAfh6i275TvudxyjosczBmL_zz9ByEUi-xXkAU,8682
 tf_keras/src/layers/attention/base_dense_attention.py,sha256=cEzBldjwQfuJfNZRimW5s-NqyENU2-lmqaNNxAGxhKw,10856
-tf_keras/src/layers/attention/multi_head_attention.py,sha256=05RC-2BSmCcBFtVY2loQPeiMYp8XArmbvovPl8kpiEA,30279
+tf_keras/src/layers/attention/multi_head_attention.py,sha256=FQX0YtXRy5kg8OlShA7cp2kfczzeWb9Oj3tbzkukLRw,30618
 tf_keras/src/layers/convolutional/__init__.py,sha256=U-4tja5JhSUva2G9uMmsZyZty2N2N9jT6EJRu5HAo-Y,3355
 tf_keras/src/layers/convolutional/base_conv.py,sha256=jvm4elEyIVSNfYZxh4inzQ1Q2CKS_f8VawvXMIJFSC4,17574
 tf_keras/src/layers/convolutional/base_depthwise_conv.py,sha256=SVgR2Y8dpeX4eDEF1e0UY0Mxh4A47eGHhJCQ1peGwNQ,9661
@@ -584,7 +584,7 @@ tf_keras/src/utils/legacy/__init__.py,sha256=EfMmeHYDzwvxNaktPhQbkTdcPSIGCqMhBND
 tf_keras/utils/__init__.py,sha256=b7_d-USe_EmLo02_P99Q1rUCzKBYayPCfiYFStP-0nw,2735
 tf_keras/utils/experimental/__init__.py,sha256=DzGogE2AosjxOVILQBT8PDDcqbWTc0wWnZRobCdpcec,97
 tf_keras/utils/legacy/__init__.py,sha256=7ujlDa5HeSRcth2NdqA0S1P2-VZF1kB3n68jye6Dj-8,189
-tf_keras_nightly-2.21.0.dev2025102309.dist-info/METADATA,sha256=nu9ClNRRGhxmwKBJPmbt7xvvCjW3WdkyJ6jv1Sm6wQE,1857
-tf_keras_nightly-2.21.0.dev2025102309.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-tf_keras_nightly-2.21.0.dev2025102309.dist-info/top_level.txt,sha256=LC8FK7zHDNKxB17C6lGKvrZ_fZZGJsRiBK23SfiDegY,9
-tf_keras_nightly-2.21.0.dev2025102309.dist-info/RECORD,,
+tf_keras_nightly-2.21.0.dev2025120310.dist-info/METADATA,sha256=lNIVZpyu96YbohxeTtctw3cPKUP1BpgWL0wovdoCbcM,1857
+tf_keras_nightly-2.21.0.dev2025120310.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+tf_keras_nightly-2.21.0.dev2025120310.dist-info/top_level.txt,sha256=LC8FK7zHDNKxB17C6lGKvrZ_fZZGJsRiBK23SfiDegY,9
+tf_keras_nightly-2.21.0.dev2025120310.dist-info/RECORD,,