tf-keras-nightly 2.19.0.dev2024121210__py3-none-any.whl → 2.21.0.dev2025123010__py3-none-any.whl

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__(
@@ -794,8 +807,7 @@ class DatasetAdapter(DataAdapter):
         # each epoch.
         return (
             self._user_steps is None
-            or tf.data.experimental.cardinality(self._dataset).numpy()
-            == self._user_steps
+            or self._dataset.cardinality().numpy() == self._user_steps
         )
 
     def _validate_args(self, y, sample_weights, steps, pss_evaluation_shards):
@@ -819,8 +831,8 @@ class DatasetAdapter(DataAdapter):
                         "specify the number of steps to run."
                     )
             else:
-                size = tf.data.experimental.cardinality(self._dataset).numpy()
-                if size == tf.data.experimental.INFINITE_CARDINALITY:
+                size = self._dataset.cardinality().numpy()
+                if size == tf.data.INFINITE_CARDINALITY:
                     if pss_evaluation_shards:
                         raise ValueError(
                             "When performing exact evaluation, the dataset "
@@ -1481,8 +1493,8 @@ class DataHandler:
         if not isinstance(dataset, tf.data.Dataset):
             return None
 
-        size = tf.data.experimental.cardinality(dataset)
-        if size == tf.data.experimental.INFINITE_CARDINALITY and steps is None:
+        size = dataset.cardinality()
+        if size == tf.data.INFINITE_CARDINALITY and steps is None:
             raise ValueError(
                 "When passing an infinitely repeating dataset, please specify "
                 "a `steps_per_epoch` value so that epoch level "
@@ -1945,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/engine/sequential.py

@@ -285,12 +285,16 @@ class Sequential(functional.Functional):
         ):
             # Determine whether the input shape is novel, i.e. whether the model
             # should be rebuilt.
-            input_shape = tuple(input_shape)
+            input_shape = tf_utils.convert_shapes(input_shape)
             if self._inferred_input_shape is None:
                 new_shape = input_shape
             else:
-                new_shape = relax_input_shape(
-                    self._inferred_input_shape, input_shape
+                new_shape = tf.nest.map_structure(
+                    _relax_input_shape,
+                    tf_utils.convert_shapes(
+                        self._inferred_input_shape, to_tuples=False
+                    ),
+                    tf_utils.convert_shapes(input_shape, to_tuples=False),
                 )
             if (
                 new_shape is not None
@@ -299,10 +303,13 @@ class Sequential(functional.Functional):
                 # A novel shape has been received: we need to rebuild the model.
                 # In case we are inside a graph function, we step out of it.
                 with tf.init_scope():
-                    inputs = input_layer.Input(
-                        batch_shape=new_shape,
-                        dtype=input_dtype,
-                        name=self.layers[0].name + "_input",
+                    inputs = tf.nest.map_structure(
+                        lambda s: input_layer.Input(
+                            batch_shape=tf_utils.convert_shapes(s),
+                            dtype=input_dtype,
+                            name=self.layers[0].name + "_input",
+                        ),
+                        tf_utils.convert_shapes(new_shape, to_tuples=False),
                     )
                     layer_input = inputs
                     created_nodes = set()
@@ -370,7 +377,7 @@ class Sequential(functional.Functional):
                 raise ValueError("You must provide an `input_shape` argument.")
             self._build_graph_network_for_inferred_shape(input_shape)
             if not self.built:
-                input_shape = tuple(input_shape)
+                input_shape = tf_utils.convert_shapes(input_shape)
                 self._build_input_shape = input_shape
                 super().build(input_shape)
         self.built = True
@@ -435,7 +442,8 @@ class Sequential(functional.Functional):
     def get_config(self):
         layer_configs = []
         serialize_obj_fn = serialization_lib.serialize_keras_object
-        if getattr(self, "use_legacy_config", None):
+        use_legacy_config = getattr(self, "use_legacy_config", False)
+        if use_legacy_config:
             serialize_obj_fn = legacy_serialization.serialize_keras_object
         for layer in super().layers:
             # `super().layers` include the InputLayer if available (it is
@@ -446,7 +454,11 @@ class Sequential(functional.Functional):
         config = training.Model.get_config(self)
         config["name"] = self.name
         config["layers"] = copy.deepcopy(layer_configs)
-        if not self._is_graph_network and self._build_input_shape is not None:
+        if (
+            use_legacy_config
+            and not self._is_graph_network
+            and self._build_input_shape
+        ):
             config["build_input_shape"] = self._build_input_shape
         return config
 
@@ -458,6 +470,7 @@ class Sequential(functional.Functional):
             layer_configs = config["layers"]
         else:
             name = None
+            build_input_shape = None
             layer_configs = config
         model = cls(name=name)
         for layer_config in layer_configs:
@@ -519,11 +532,15 @@ def _get_shape_tuple(t):
     return None
 
 
-def relax_input_shape(shape_1, shape_2):
+def _relax_input_shape(shape_1, shape_2):
     if shape_1 is None or shape_2 is None:
         return None
-    if len(shape_1) != len(shape_2):
+    if shape_1.rank is None or shape_2.rank is None:
+        return None
+    if shape_1.rank != shape_2.rank:
         return None
+    shape_1 = shape_1.as_list()
+    shape_2 = shape_2.as_list()
     return tuple(None if d1 != d2 else d1 for d1, d2 in zip(shape_1, shape_2))
 
 

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/src/layers/core/tf_op_layer.py

@@ -259,6 +259,10 @@ class TFOpLambda(Layer):
 
         self._call_spec.expects_training_arg = False
         self._call_spec.expects_mask_arg = False
+        # Clear the call-context arguments for the layer's call method.
+        # Otherwise, Keras ends up injecting context arguments into the op-call
+        # when the call method accepts kwargs.
+        self._call_spec._expected_context_args.clear()
 
     def _call_wrapper(self, *args, **kwargs):
         created_variables = []

tf_keras/src/layers/normalization/spectral_normalization.py

@@ -95,7 +95,7 @@ class SpectralNormalization(Wrapper):
 
     def call(self, inputs, training=False):
         if training:
-            self.normalize_weights()
+            self._update_weights()
 
         output = self.layer(inputs)
         return output
@@ -105,35 +105,42 @@ class SpectralNormalization(Wrapper):
             self.layer.compute_output_shape(input_shape).as_list()
         )
 
+    def _update_weights(self):
+        weights = self.kernel
+        vector_u = self.vector_u
+
+        kernel_weights, vector_u = tf.cond(
+            tf.reduce_all(tf.equal(weights, 0)),
+            lambda: (weights, vector_u),
+            lambda: self.normalize_weights(),
+        )
+        self.kernel.assign(kernel_weights)
+        self.vector_u.assign(vector_u)
+
     def normalize_weights(self):
         """Generate spectral normalized weights.
 
         This method will update the value of `self.kernel` with the
         spectral normalized value, so that the layer is ready for `call()`.
         """
-
-        weights = tf.reshape(self.kernel, [-1, self.kernel_shape[-1]])
+        # Initialize vector_v to hint the compiler it always exist.
         vector_u = self.vector_u
-
-        # check for zeroes weights
-        if not tf.reduce_all(tf.equal(weights, 0.0)):
-            for _ in range(self.power_iterations):
-                vector_v = tf.math.l2_normalize(
-                    tf.matmul(vector_u, weights, transpose_b=True)
-                )
-                vector_u = tf.math.l2_normalize(tf.matmul(vector_v, weights))
-            vector_u = tf.stop_gradient(vector_u)
-            vector_v = tf.stop_gradient(vector_v)
-            sigma = tf.matmul(
-                tf.matmul(vector_v, weights), vector_u, transpose_b=True
-            )
-            self.vector_u.assign(tf.cast(vector_u, self.vector_u.dtype))
-            self.kernel.assign(
-                tf.cast(
-                    tf.reshape(self.kernel / sigma, self.kernel_shape),
-                    self.kernel.dtype,
-                )
+        vector_v = self.vector_u
+        weights = tf.reshape(self.kernel, [-1, self.kernel_shape[-1]])
+        for _ in range(self.power_iterations):
+            vector_v = tf.math.l2_normalize(
+                tf.matmul(vector_u, weights, transpose_b=True)
             )
+            vector_u = tf.math.l2_normalize(tf.matmul(vector_v, weights))
+        vector_u = tf.stop_gradient(vector_u)
+        vector_v = tf.stop_gradient(vector_v)
+        sigma = tf.matmul(
+            tf.matmul(vector_v, weights),
+            vector_u,
+            transpose_b=True,
+        )
+        weights_normalized = tf.reshape(weights / sigma, self.kernel_shape)
+        return weights_normalized, vector_u
 
     def get_config(self):
         config = {"power_iterations": self.power_iterations}

tf_keras/src/layers/rnn/cell_wrappers.py

@@ -52,9 +52,21 @@ class _RNNCellWrapper(AbstractRNNCell):
         super().__init__(*args, **kwargs)
         self.cell = cell
         cell_call_spec = tf_inspect.getfullargspec(cell.call)
+        accepts_kwargs = cell_call_spec.varkw is not None
+
         self._call_spec.expects_training_arg = (
             "training" in cell_call_spec.args
-        ) or (cell_call_spec.varkw is not None)
+        ) or accepts_kwargs
+
+        # Filter _expects_context_arg. An argument is kept if:
+        # 1. It's an explicit argument in cell_call_spec.args OR
+        # 2. The cell accepts arbitrary keyword arguments (**kwargs),
+        #    meaning it could potentially handle the context argument.
+        self._call_spec._expected_context_args = {
+            arg
+            for arg in self._call_spec._expected_context_args
+            if (arg in cell_call_spec.args) or accepts_kwargs
+        }
 
     def _call_wrapped_cell(self, inputs, state, cell_call_fn, **kwargs):
         """Calls the wrapped cell and performs the wrapping logic.