tf-keras-nightly 2.20.0.dev2025052009__py3-none-any.whl → 2.20.0.dev2025052209__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.20.0.dev2025052009"
+__version__ = "2.20.0.dev2025052209"

tf_keras/src/engine/base_layer.py

@@ -308,6 +308,7 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
         self, trainable=True, name=None, dtype=None, dynamic=False, **kwargs
     ):
         self._instrument_layer_creation()
+        self._called = False
 
         # These properties should be set by the user via keyword arguments.
         # note that 'dtype', 'input_shape' and 'batch_input_shape'
@@ -326,6 +327,10 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
         # Validate optional keyword arguments.
         generic_utils.validate_kwargs(kwargs, allowed_kwargs)
 
+        # Track the built-in call-context arguments. These are arguments that
+        # are tracked and propagated across the call-stack by default.
+        self._call_context_args = {"training"}
+
         # Mutable properties
         # Indicates whether the layer's weights are updated during training
         # and whether the layer's updates are run during training.
@@ -411,6 +416,9 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
 
         self._init_call_fn_args()
 
+        # Track the built-in call-context arguments.
+        self._call_spec._update_call_context_arguments(self._call_context_args)
+
         # Whether the `call` method can be used to build a TF graph without
         # issues.  This attribute has no effect if the model is created using
         # the Functional API. Instead, `model.dynamic` is determined based on
@@ -1042,6 +1050,7 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
         # - input_spec compatibility is only checked against `inputs`
         # - mixed precision casting (autocast) is only applied to `inputs`,
         #   not to any other argument.
+        self._called = True
         inputs, args, kwargs = self._call_spec.split_out_first_arg(args, kwargs)
         input_list = tf.nest.flatten(inputs)
 
@@ -1080,17 +1089,21 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
         if self._expects_mask_arg and mask_is_implicit:
             kwargs["mask"] = input_masks
 
-        # Training mode for `Layer.call` is set via (in order of priority):
-        # (1) The `training` argument passed to this `Layer.call`, if it is not
-        #  None
-        # (2) The training mode of an outer `Layer.call`.
-        # (3) The default mode set by `tf.keras.backend.set_learning_phase` (if
-        #  set)
-        # (4) Any non-None default value for `training` specified in the call
+        # Call-context arguments for `Layer.call` is set via (in order of
+        # priority):
+        # (1) The argument passed to this `Layer.call`, if it is not None
+        # (2) The argument value of an outer `Layer.call`.
+        # (3) (only for "training") The default mode set by
+        # `tf.keras.backend.set_learning_phase` (if set)
+        # (4) Any non-None default value for the argument specified in the call
         #  signature
-        # (5) False (treating the layer as if it's in inference)
-        args, kwargs, training_mode = self._set_training_mode(
-            args, kwargs, call_context
+        # (5) False
+        (
+            args,
+            kwargs,
+            propagated,
+        ) = self._get_propagated_call_context_arguments(
+            args, kwargs, call_context, self._call_context_args
         )
 
         # Losses are cleared for all sublayers on the outermost `Layer.call`.
@@ -1104,7 +1117,7 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
             layer=self,
             inputs=inputs,
             build_graph=not eager,
-            training=training_mode,
+            call_context_args=propagated,
         ):
             input_spec.assert_input_compatibility(
                 self.input_spec, inputs, self.name
@@ -1152,6 +1165,55 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
 
                 return outputs
 
+    def _register_call_context_args(self, *argument_names):
+        """Registers call-context args for this layer.
+        If this layer declares a `call()` method that accepts
+        one or more of the given args, those args will be
+        automatically injected into the call signature of this
+        layer. This layer will also propagate the args to any
+        nested sublayers that are called from within this layer.
+        If this layer doesn't declare a `call()` method that
+        accepts one or more of the given args, these args will
+        simply be propagated to any nested sublayers without
+        being injected into the call signature of this layer.
+        This is useful for propagating custom arguments
+        from top-level layers/models to sublayers.
+        Example:
+        ```
+        class Inner(layers.Layer):
+            def __init__(self):
+                super().__init__()
+                # Register `foo_mode` as a call-context arg
+                self._register_call_context_args("foo_mode")
+            def call(self, x, foo_mode=False):
+                # If foo_mode=True add 1, otherwise add 0
+                add_val = ops.where(foo_mode, 1.0, 0.0)
+                return x + add_val
+        class Outer(layers.Layer):
+            def __init__(self):
+                super().__init__()
+                self.inner = Inner()
+            def call(self, x):
+                # We don't explicitly pass foo_mode here—Base Layer.__call__
+                # should inject it into `self.inner`
+                return self.inner(x)
+        sample_input = np.array([[1.0], [2.0]])
+        # Sequential model
+        seq = models.Sequential([Outer()])
+        # Tell the Sequential model to propagate foo_mode down
+        # the call-stack
+        seq.register_call_context_args("foo_mode")
+        # foo_mode=True -> input + 1
+        out_true = seq(sample_input, foo_mode=True)
+        """
+        if self._called:
+            raise RuntimeError(
+                "Cannot add call-context args after the layer has been called."
+            )
+        self._call_context_args |= set(argument_names)
+        self._call_spec._update_call_context_arguments(argument_names)
+        self._call_spec._update_call_context_argument_defaults(argument_names)
+
     def _get_unnested_name_scope(self):
         if _is_name_scope_on_model_declaration_enabled:
             with _name_scope_unnester(
@@ -2535,47 +2597,57 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
             kwargs["mask"] = input_masks
             mask_arg_passed_by_framework = True
 
-        # If `training` argument is None or not explicitly passed,
-        # propagate `training` value from this layer's calling layer.
-        training_value = None
-        training_arg_passed_by_framework = False
-        # Priority 1: `training` was explicitly passed a non-None value.
-        if self._call_spec.arg_was_passed("training", args, kwargs):
-            training_value = self._call_spec.get_arg_value(
-                "training", args, kwargs
-            )
-            if not self._expects_training_arg:
-                kwargs.pop("training")
-
-        if training_value is None:
-            # Priority 2: `training` was passed to a parent layer.
-            if call_context.training is not None:
-                training_value = call_context.training
-            # Priority 3: `learning_phase()` has been set.
-            elif backend.global_learning_phase_is_set():
-                training_value = backend.learning_phase()
-                # Force the training_value to be bool type which matches to the
-                # contract for layer/model call args.
-                if tf.is_tensor(training_value):
-                    training_value = tf.cast(training_value, tf.bool)
+        propagated = dict()
+        args_passed_by_framework = dict()
+        for context_arg in self._call_context_args:
+            # If `training` argument is None or not explicitly passed,
+            # propagate `training` value from this layer's calling layer.
+            value = None
+            args_passed_by_framework[context_arg] = False
+            # Priority 1: `training` was explicitly passed a non-None value.
+            if self._call_spec.arg_was_passed(context_arg, args, kwargs):
+                value = self._call_spec.get_arg_value(context_arg, args, kwargs)
+                if not self._expects_context_arg(context_arg):
+                    kwargs.pop(context_arg)
+
+            if value is None:
+                # Priority 2: `training` was passed to a parent layer.
+                if call_context.get_call_context_arg(context_arg) is not None:
+                    value = call_context.get_call_context_arg(context_arg)
+                # Priority 3: `learning_phase()` has been set.
+                elif (
+                    context_arg == "training"
+                    and backend.global_learning_phase_is_set()
+                ):
+                    value = backend.learning_phase()
+                    # Force the training_value to be bool type which matches to
+                    # the contract for layer/model call args.
+                    if tf.is_tensor(value):
+                        value = tf.cast(value, tf.bool)
+                    else:
+                        value = bool(value)
+                # Priority 4: trace layer with the default training argument
+                # specified in the `call` signature (or in inference mode if the
+                # `call` signature specifies no non-None default).
                 else:
-                    training_value = bool(training_value)
-            # Priority 4: trace layer with the default training argument
-            # specified in the `call` signature (or in inference mode if the
-            # `call` signature specifies no non-None default).
-            else:
-                training_value = self._call_spec.default_training_arg
-            # In cases (2), (3), (4) the training argument is passed
-            # automatically by the framework, and will not be hard-coded into
-            # the model.
-            if self._expects_training_arg:
-                args, kwargs = self._call_spec.set_arg_value(
-                    "training", training_value, args, kwargs
-                )
-                training_arg_passed_by_framework = True
+                    value = self._call_spec.get_context_arg_default(context_arg)
+                # In cases (2), (3), (4) the training argument is passed
+                # automatically by the framework, and will not be hard-coded
+                # into the model.
+                if self._expects_context_arg(context_arg):
+                    args, kwargs = self._call_spec.set_arg_value(
+                        context_arg, value, args, kwargs
+                    )
+                    args_passed_by_framework[context_arg] = True
+
+            if value is not None:
+                propagated[context_arg] = value
 
         with call_context.enter(
-            layer=self, inputs=inputs, build_graph=True, training=training_value
+            layer=self,
+            inputs=inputs,
+            build_graph=True,
+            call_context_args=propagated,
         ):
             # Check input assumptions set after layer building, e.g. input
             # shape.
@@ -2601,10 +2673,13 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
                     "Tensor or a list of Tensors, not None "
                     "(layer: " + self.name + ")."
                 )
-            if training_arg_passed_by_framework:
-                args, kwargs = self._call_spec.set_arg_value(
-                    "training", None, args, kwargs, pop_kwarg_if_none=True
-                )
+
+            for context_arg, is_passed in args_passed_by_framework.items():
+                if is_passed:
+                    args, kwargs = self._call_spec.set_arg_value(
+                        context_arg, None, args, kwargs, pop_kwarg_if_none=True
+                    )
+
             if mask_arg_passed_by_framework:
                 kwargs.pop("mask")
             # Node connectivity does not special-case the first argument.
@@ -2613,52 +2688,100 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
             )
             return outputs
 
-    def _set_training_mode(self, args, kwargs, call_context):
-        training_mode = None
-        if self._expects_training_arg:
-            # (1) `training` was passed to this `Layer.call`.
-            if self._call_spec.arg_was_passed("training", args, kwargs):
-                training_mode = self._call_spec.get_arg_value(
-                    "training", args, kwargs
-                )
-            # If no `training` arg was passed, or `None` was explicitly passed,
-            # the framework will make a decision about the training mode is.
-            if training_mode is None:
-                call_ctx_training = call_context.training
-                # (2) `training` mode is inferred from an outer `Layer.call`.
-                if call_ctx_training is not None:
-                    training_mode = call_ctx_training
-                # (3) User set `tf.keras.backend.set_learning_phase`.
-                elif backend.global_learning_phase_is_set():
-                    training_mode = backend.learning_phase()
-                    # Ensure value is a `bool` or `tf.bool`.
-                    if isinstance(training_mode, bool):
-                        pass
-                    elif tf.is_tensor(training_mode):
-                        training_mode = tf.cast(training_mode, tf.bool)
+    def _get_propagated_call_context_arguments(
+        self, args, kwargs, call_context, local_call_context_arguments
+    ):
+        """Resolves the values for propagated call context arguments for the
+        current layer.
+
+        Args:
+            args: The arguments passed to the current layer's `call` method.
+            kwargs: The keyword arguments passed to the current layer's `call`
+              method.
+            call_context: The `CallContext` for the current call-stack.
+            local_call_context_arguments: The call-context arguments registered
+            with to the current layer's `Layer.call` method.
+
+        Returns:
+            A tuple of the following:
+            1. Updated args
+            2. Updated kwargs
+            3. A dictionary of the resolved call-context arguments that should
+            be propagated to the next layer in the call-stack.
+        """
+        propagated_context = dict()
+        relevant_arguments = call_context.call_context_args.keys() | set(
+            local_call_context_arguments
+        )
+
+        for argument in relevant_arguments:
+            authoritative_value = None
+            was_explicitly_passed = self._call_spec.arg_was_passed(
+                argument, args, kwargs
+            )
+            if self._expects_context_arg(argument):
+                # (1) `arg_name` was passed to this `Layer.call`.
+                if was_explicitly_passed:
+                    authoritative_value = self._call_spec.get_arg_value(
+                        argument, args, kwargs
+                    )
+                # If no `arg_name` arg was passed, or `None` was explicitly
+                # passed, the framework will make a decision about the training
+                # mode is.
+                if authoritative_value is None:
+                    value_from_context = call_context.get_call_context_arg(
+                        argument
+                    )
+                    # (2) `arg_name` mode is inferred from an outer
+                    # `Layer.call`.
+                    if value_from_context is not None:
+                        authoritative_value = value_from_context
+                    # (3) User set `tf.keras.backend.set_learning_phase`.
+                    elif (
+                        argument == "training"
+                        and backend.global_learning_phase_is_set()
+                    ):
+                        authoritative_value = backend.learning_phase()
+                        # Ensure value is a `bool` or `tf.bool`.
+                        if isinstance(authoritative_value, bool):
+                            pass
+                        elif tf.is_tensor(authoritative_value):
+                            authoritative_value = tf.cast(
+                                authoritative_value, tf.bool
+                            )
+                        else:
+                            authoritative_value = bool(authoritative_value)
+                    # (4) We default to using `call`'s default value for
+                    # `arg_name`, or treating the layer as if it is in inference
+                    # if no non-None default is specified in the `call`
+                    # signature.
                     else:
-                        training_mode = bool(training_mode)
-                # (4) We default to using `call`'s default value for `training`,
-                # or treating the layer as if it is in inference if no non-None
-                # default is specified in the `call` signature.
-                else:
-                    training_mode = self._call_spec.default_training_arg
+                        authoritative_value = (
+                            self._call_spec.get_context_arg_default(argument)
+                        )
 
-                # For case (2), (3), (4) `training` arg is passed by framework.
-                args, kwargs = self._call_spec.set_arg_value(
-                    "training", training_mode, args, kwargs
-                )
-        else:
-            if "training" in kwargs:
-                # `training` was passed to this `Layer` but is not needed for
-                # `Layer.call`. It will set the default mode for inner
-                # `Layer.call`s.
-                training_mode = kwargs.pop("training")
+                    # For case (2), (3), (4) `arg_name` arg is passed by
+                    # framework.
+                    args, kwargs = self._call_spec.set_arg_value(
+                        argument, authoritative_value, args, kwargs
+                    )
             else:
-                # Grab the current `training` mode from any outer `Layer.call`.
-                training_mode = call_context.training
+                if argument in kwargs:
+                    # `arg_name` was passed to this `Layer` but is not needed
+                    # for `Layer.call`. It will set the default mode for inner
+                    # `Layer.call`s.
+                    authoritative_value = kwargs.pop(argument)
+                else:
+                    # Grab the current `arg_name` mode from any outer
+                    # `Layer.call`.
+                    authoritative_value = call_context.get_call_context_arg(
+                        argument
+                    )
 
-        return args, kwargs, training_mode
+            if authoritative_value is not None:
+                propagated_context[argument] = authoritative_value
+
+        return args, kwargs, propagated_context
 
     def _autographed_call(self):
         # Wrapping `call` function in autograph to allow for dynamic control
@@ -3351,7 +3474,8 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
 
     def _init_call_fn_args(self, expects_training_arg=None):
         self._call_spec = layer_utils.CallFunctionSpec(
-            tf_inspect.getfullargspec(self.call)
+            tf_inspect.getfullargspec(self.call),
+            getattr(self, "_call_context_args", set()),
         )
         if expects_training_arg is not None:
             self._call_spec.expects_training_arg = expects_training_arg
@@ -3361,6 +3485,9 @@ class Layer(tf.Module, version_utils.LayerVersionSelector):
         """Whether the call function uses 'training' as a parameter."""
         return self._call_spec.expects_training_arg
 
+    def _expects_context_arg(self, argument_name):
+        return argument_name in self._call_spec.expected_context_args
+
     @property
     def _expects_mask_arg(self):
         return self._call_spec.expects_mask_arg

tf_keras/src/engine/base_layer_utils.py

@@ -480,7 +480,8 @@ class CallContext:
       layer: The `Layer` whose `call` is currently active.
       inputs: The inputs to the currently active `Layer`.
       build_graph: Whether currently inside a Graph or FuncGraph.
-      training: Whether currently executing in training or inference mode.
+      call_context_args: The call-context arguments being propagated through the
+        the call-stack.
       saving: Whether currently saving to SavedModel.
       frozen: Whether currently executing inside a `Layer` with `trainable` set
         to `False`.
@@ -495,6 +496,7 @@ class CallContext:
             "layer": None,
             "inputs": None,
             "build_graph": False,
+            "call_context_args": dict(),
             "training": None,
             "saving": None,
         }
@@ -502,14 +504,17 @@ class CallContext:
         # refactor.
         self._in_keras_graph = False
 
-    def enter(self, layer, inputs, build_graph, training, saving=None):
+    def enter(
+        self, layer, inputs, build_graph, call_context_args=dict(), saving=None
+    ):
         """Push a Layer and its inputs and state onto the current call context.
 
         Args:
           layer: The `Layer` whose `call` is currently active.
           inputs: The inputs to the currently active `Layer`.
           build_graph: Whether currently inside a Graph or FuncGraph.
-          training: Whether currently executing in training or inference mode.
+          call_context_args: The call-context arguments being propagated through
+            the call-stack.
           saving: Whether currently saving to SavedModel.
 
         Returns:
@@ -519,7 +524,7 @@ class CallContext:
             "layer": layer,
             "inputs": inputs,
             "build_graph": build_graph,
-            "training": training,
+            "call_context_args": call_context_args,
             "saving": saving,
         }
         return CallContextManager(self, state)
@@ -538,7 +543,14 @@ class CallContext:
 
     @property
     def training(self):
-        return self._state["training"]
+        return self.call_context_args.get("training", None)
+
+    @property
+    def call_context_args(self):
+        return self._state["call_context_args"]
+
+    def get_call_context_arg(self, arg_name):
+        return self.call_context_args.get(arg_name, None)
 
     @property
     def saving(self):

tf_keras/src/engine/base_layer_v1.py

@@ -132,6 +132,7 @@ class Layer(base_layer.Layer):
         self, trainable=True, name=None, dtype=None, dynamic=False, **kwargs
     ):
         self._instrument_layer_creation()
+        self._called = False
 
         # These properties should be set by the user via keyword arguments.
         # note that 'dtype', 'input_shape' and 'batch_input_shape'
@@ -165,6 +166,8 @@ class Layer(base_layer.Layer):
         self._input_spec = None
         self.supports_masking = False
 
+        self._call_context_args = {"training"}
+
         self._init_set_name(name)
         self._activity_regularizer = regularizers.get(
             kwargs.pop("activity_regularizer", None)
@@ -705,6 +708,7 @@ class Layer(base_layer.Layer):
           RuntimeError: if `super().__init__()` was not called in the
             constructor.
         """
+        self._called = True
         self._assert_built_as_v1()
 
         if not hasattr(self, "_thread_local"):
@@ -803,7 +807,12 @@ class Layer(base_layer.Layer):
         if build_graph and base_layer_utils.needs_keras_history(inputs):
             base_layer_utils.create_keras_history(inputs)
 
-        with call_context.enter(self, inputs, build_graph, training_value):
+        with call_context.enter(
+            self,
+            inputs,
+            build_graph,
+            call_context_args={"training": training_value},
+        ):
             # Check input assumptions set after layer building, e.g. input
             # shape.
             if build_graph:

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/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.

tf_keras/src/saving/legacy/saved_model/save_impl.py

@@ -219,7 +219,11 @@ def wrap_layer_functions(layer, serialization_cache):
     with tracing_scope():
         call_collection.trace_with_input_signature()
         with base_layer_utils.call_context().enter(
-            layer, inputs=None, build_graph=True, training=None, saving=True
+            layer,
+            inputs=None,
+            build_graph=True,
+            call_context_args={},
+            saving=True,
         ):
             for fn in fns.values():
                 if fn is not None and not isinstance(fn, LayerCall):
@@ -515,19 +519,28 @@ class LayerCallCollection:
             else:
                 add_trace_to_queue(fn, args, kwargs)
 
-    def training_arg_was_passed(self, args, kwargs):
+    def arg_was_passed(self, arg_name, args, kwargs):
+        """Returns True if the argument was passed to the call function."""
         return self._call_spec.arg_was_passed(
-            "training", args, kwargs, inputs_in_args=True
+            arg_name, args, kwargs, inputs_in_args=True
         )
 
-    def get_training_arg_value(self, args, kwargs):
+    def training_arg_was_passed(self, args, kwargs):
+        """Returns True if the training arg was passed to the call function."""
+        return self.arg_was_passed("training", args, kwargs)
+
+    def get_arg_value(self, arg_name, args, kwargs):
+        """Returns the value of the given argument or None if not found."""
         try:
             return self._call_spec.get_arg_value(
-                "training", args, kwargs, inputs_in_args=True
+                arg_name, args, kwargs, inputs_in_args=True
             )
-        except KeyError:  # Training is not in args or kwargs.
+        except KeyError:  # Arg not found in args or kwargs.
             return None
 
+    def get_training_arg_value(self, args, kwargs):
+        return self.get_arg_value("training", args, kwargs)
+
     def get_input_arg_value(self, args, kwargs):
         return self._call_spec.get_arg_value(
             self._input_arg_name, args, kwargs, inputs_in_args=True
@@ -613,20 +626,23 @@ def layer_call_wrapper(call_collection, method, name):
     def wrapper(*args, **kwargs):
         """Calls method within call context."""
         layer = call_collection.layer
-        training = None
+        propagated = {"training": None}
         inputs = _filtered_inputs([args, kwargs])
 
-        if (args or kwargs) and call_collection.training_arg_was_passed(
-            args, kwargs
-        ):
-            training = call_collection.get_training_arg_value(args, kwargs)
+        for context_arg in layer._call_context_args:
+            if (args or kwargs) and call_collection.arg_was_passed(
+                context_arg, args, kwargs
+            ):
+                propagated[context_arg] = call_collection.get_arg_value(
+                    context_arg, args, kwargs
+                )
 
         original_losses = _reset_layer_losses(layer)
         with base_layer_utils.call_context().enter(
             layer,
             inputs=inputs,
             build_graph=False,
-            training=training,
+            call_context_args=propagated,
             saving=True,
         ):
             with autocast_variable.enable_auto_cast_variables(

tf_keras/src/saving/legacy/saving_utils.py

@@ -138,12 +138,24 @@ def trace_model_call(model, input_signature=None):
     @tf.function
     def _wrapped_model(*args, **kwargs):
         """A concrete tf.function that wraps the model's call function."""
+        call_context = base_layer_utils.call_context()
+
+        args, kwargs, propagated = model._get_propagated_call_context_arguments(
+            args, kwargs, call_context, model._call_context_args
+        )
+
         (args, kwargs,) = model._call_spec.set_arg_value(
             "training", False, args, kwargs, inputs_in_args=True
         )
 
-        with base_layer_utils.call_context().enter(
-            model, inputs=None, build_graph=False, training=False, saving=True
+        propagated["training"] = False
+
+        with call_context.enter(
+            model,
+            inputs=None,
+            build_graph=False,
+            call_context_args=propagated,
+            saving=True,
         ):
             outputs = model(*args, **kwargs)
 

tf_keras/src/utils/layer_utils.py

@@ -775,11 +775,13 @@ class CallFunctionSpec:
     """Caches the spec and provides utilities for handling call function
     args."""
 
-    def __init__(self, full_argspec):
+    def __init__(self, full_argspec, call_context_args=set()):
         """Initialies a `CallFunctionSpec`.
 
         Args:
           full_argspec: the FullArgSpec of a call function of a layer.
+          call_context_args: The set of call-context arguments registered
+            with to the current layer.
         """
         self._full_argspec = full_argspec
 
@@ -797,6 +799,18 @@ class CallFunctionSpec:
             "mask" in self._arg_names or call_accepts_kwargs
         )
 
+        # Track the set of call-context arguments that the current layer's
+        # `call` method accepts.
+        self._expected_context_args = set()
+        self._update_call_context_arguments(call_context_args)
+
+        self._context_arg_defaults = dict()
+        self._update_call_context_argument_defaults(call_context_args)
+
+    def _update_call_context_argument_defaults(self, context_args):
+        """Updates the set of call-context argument defaults for the current
+        layer's `call` method.
+        """
         call_fn_defaults = self._full_argspec.defaults or []
         defaults = dict()
         # The call arg defaults are an n-tuple of the last n elements of the
@@ -806,7 +820,21 @@ class CallFunctionSpec:
         # The default training arg will be any (non-None) default specified in
         # the method signature, or None if no value is specified.
         defaults.update(self._full_argspec.kwonlydefaults or {})
-        self._default_training_arg = defaults.get("training")
+
+        for arg in context_args:
+            self._context_arg_defaults[arg] = defaults.get(arg)
+
+    def _update_call_context_arguments(self, context_args):
+        """Updates the set of call-context arguments that the current layer's
+        `call` method accepts.
+        """
+        call_accepts_kwargs = self._full_argspec.varkw is not None
+        args_to_add = {
+            arg
+            for arg in context_args
+            if call_accepts_kwargs or arg in self._arg_names
+        }
+        self._expected_context_args.update(args_to_add)
 
     @property
     def full_argspec(self):
@@ -843,6 +871,16 @@ class CallFunctionSpec:
     def expects_training_arg(self, value):
         self._expects_training_arg = value
 
+    @property
+    def expected_context_args(self):
+        """The set of call-context arguments that the current layer's
+        `call` method accepts."""
+        return self._expected_context_args
+
+    @expected_context_args.setter
+    def expected_context_args(self, value):
+        self._expected_context_args = value
+
     @property
     def expects_mask_arg(self):
         """Whether the call function uses `mask` as a parameter."""
@@ -855,7 +893,11 @@ class CallFunctionSpec:
     @property
     def default_training_arg(self):
         """The default value given to the "training" argument."""
-        return self._default_training_arg
+        return self.get_context_arg_default("training")
+
+    def get_context_arg_default(self, arg_name):
+        """The default value given to the call context arguments."""
+        return self._context_arg_defaults.get(arg_name, None)
 
     def arg_was_passed(self, arg_name, args, kwargs, inputs_in_args=False):
         """Returns true if argument is present in `args` or `kwargs`.

{tf_keras_nightly-2.20.0.dev2025052009.dist-info → tf_keras_nightly-2.20.0.dev2025052209.dist-info}/METADATA

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

{tf_keras_nightly-2.20.0.dev2025052009.dist-info → tf_keras_nightly-2.20.0.dev2025052209.dist-info}/RECORD

@@ -1,4 +1,4 @@
-tf_keras/__init__.py,sha256=bXm9QQlbCSpysAH2JS8-HXtaXBiDffnpJfbJLHw-Guc,911
+tf_keras/__init__.py,sha256=akRyIApBb0vOJtY9Wv6obInv-g0GJtvlL46olfFi_Rk,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
@@ -271,9 +271,9 @@ tf_keras/src/dtensor/lazy_variable.py,sha256=c3yylbga0se3Geflutss3fz5RzBYuY2vkU3
 tf_keras/src/dtensor/test_util.py,sha256=9QAbt44mlirdqwG2ertTsoXNKG2V4Z0bqJFxGdxy5BY,4572
 tf_keras/src/dtensor/utils.py,sha256=2TTSCEOA61Ia1FAPfQWJ2CRfiocBGUZreXH9UBFzFbk,6441
 tf_keras/src/engine/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-tf_keras/src/engine/base_layer.py,sha256=H7TK3ezXORn1B79TZO3klFGE_hmXHAauS3h5k_8xJvA,156553
-tf_keras/src/engine/base_layer_utils.py,sha256=YMJF5sZJhFF_yzfqOtqi4YTsyUE2ZQ_cJJOIdXnuS2w,35795
-tf_keras/src/engine/base_layer_v1.py,sha256=cX-OCSNio3Tr2M6twr_PUgKulePLZDNW_4xLXjgYbN4,102700
+tf_keras/src/engine/base_layer.py,sha256=cCSNR0ZMZBPRzpq6bByylNm5rwA8KzvIf16vKaXMpjE,161974
+tf_keras/src/engine/base_layer_utils.py,sha256=AFjqwXM-WShf0dfsyIotlXYIRJlqYyjQhAf50xZgyos,36166
+tf_keras/src/engine/base_layer_v1.py,sha256=V3Jyip5kiDyN-lwHAThS368iKxpznSw4tj4mwfWLobc,102896
 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
@@ -343,7 +343,7 @@ tf_keras/src/layers/core/embedding.py,sha256=iOdkBiP1IzwOVPjsKWA54NXrlk5KgJ0DfQ8
 tf_keras/src/layers/core/identity.py,sha256=yj5cWlUTlYq_J_ZQb1iLzM0bqaM4V6TXVwM4iuBFp9U,1301
 tf_keras/src/layers/core/lambda_layer.py,sha256=QzetX-lV9ybonQKg_6QzSm8w9Vkq8CPAM4BcAke7CZk,16481
 tf_keras/src/layers/core/masking.py,sha256=19p6HYGlKdUfQnelsAoee6wf87fWx67NSGinyjagNc4,3340
-tf_keras/src/layers/core/tf_op_layer.py,sha256=4WDRrT8dVwnD7avcWvMCk9mnGwfHcaN3Dmhf7CBeqzQ,21066
+tf_keras/src/layers/core/tf_op_layer.py,sha256=R6dFECVkPbmKi1nQVcxJy5lNxSVwiMlaWXB7j0PjI7Q,21320
 tf_keras/src/layers/experimental/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 tf_keras/src/layers/experimental/dynamic_embedding.py,sha256=KuVIawm3avPEa5c2IDOyBH14xiU5bYbPqcm_HugfWYA,10730
 tf_keras/src/layers/experimental/dynamic_lookup.py,sha256=CMNOaxAIkB1ChPcusuymhLAYTvobEbCBli6YkuWw8RE,13720
@@ -454,7 +454,7 @@ tf_keras/src/layers/rnn/base_cudnn_rnn.py,sha256=cuPVg6r4L1pVWYTp3WFbJhikuIR2Vmg
 tf_keras/src/layers/rnn/base_rnn.py,sha256=I7mWl4KQC26gILDt9pZ9moZ81yM57lvci6hzJ9ROrxo,41968
 tf_keras/src/layers/rnn/base_wrapper.py,sha256=x4GANiXtmh9ztAFh7QtfbnQE76UVCGpaHp_XhrSs0Os,3159
 tf_keras/src/layers/rnn/bidirectional.py,sha256=JyZuBU0q2lt4augThwm8vyTvYwEJxyawsHmgNIul5vU,22670
-tf_keras/src/layers/rnn/cell_wrappers.py,sha256=T3FIiY9vIr0Or1N_SWNVnHR3LH6xnJ5DgNNYLk-sV6c,26874
+tf_keras/src/layers/rnn/cell_wrappers.py,sha256=fMGpdFFoRWRIuKz88NcnMvAtevv8OYHzxkF86Ltmwfk,27384
 tf_keras/src/layers/rnn/conv_lstm1d.py,sha256=suShze6ipNXabGlKJTxkOia17ZP4SeEei3Mi4F8lFOQ,8761
 tf_keras/src/layers/rnn/conv_lstm2d.py,sha256=myxOioB3yNn0L_-gMh0R41sb-MwTXO993lAT05_N0Zw,8874
 tf_keras/src/layers/rnn/conv_lstm3d.py,sha256=GT4OoPFtCr5xgaaqy3ezt5DyDu8Ut-wQEihCOHFk0D4,8969
@@ -544,7 +544,7 @@ tf_keras/src/saving/legacy/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74
 tf_keras/src/saving/legacy/hdf5_format.py,sha256=IqFXHN96fuqKwu_akaqTyf9ISRPavP3Ahjydat948O4,42438
 tf_keras/src/saving/legacy/model_config.py,sha256=ZE6H_dKdmo2dlBWkr2nYO8SXcMEhshgza3sHPCpeu-k,4140
 tf_keras/src/saving/legacy/save.py,sha256=TdjiEamZ8MAsPAWsYMEtrdCRppbHcBIwJh9eVfdUS3k,23612
-tf_keras/src/saving/legacy/saving_utils.py,sha256=VTxnYFSWZ7m_40deANuWbtyhyXq0o0c5vLuBeCbgwi8,13745
+tf_keras/src/saving/legacy/saving_utils.py,sha256=0iXchqZQNw9s5kB9_7SIj2p3Qd_21jdfHJQ3b3YVWQs,14042
 tf_keras/src/saving/legacy/serialization.py,sha256=OrmHQPolQFsR-UCMxNTxkIFTKY4DKcAgMm1jdhF7TqU,22285
 tf_keras/src/saving/legacy/saved_model/__init__.py,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
 tf_keras/src/saving/legacy/saved_model/base_serialization.py,sha256=dALR19_zt4c80zVw3yjCj9wfRoJufDjCrvkJyS82Dnk,5104
@@ -558,7 +558,7 @@ tf_keras/src/saving/legacy/saved_model/model_serialization.py,sha256=IxQ1TfBGagV
 tf_keras/src/saving/legacy/saved_model/network_serialization.py,sha256=ofbKN9V3syw0AQebgy2PlvaiAHi3SnBFTg-PUgclTng,1180
 tf_keras/src/saving/legacy/saved_model/order_preserving_set.py,sha256=zvNFzss8wSc0vngv74dNnQO_hxpxmEWWBBv1TTLsbPY,3250
 tf_keras/src/saving/legacy/saved_model/save.py,sha256=2-AaGFhFxzfZLkIW1qx9-rTcaZvYMFkQYP7ijfwA-ZI,6395
-tf_keras/src/saving/legacy/saved_model/save_impl.py,sha256=cWDBJ0uYfV1dNB5pdKFqDwYLQaaxWNeG8GSLiABAmec,29751
+tf_keras/src/saving/legacy/saved_model/save_impl.py,sha256=mcdNPwJYwzOsdSisgwkBEbnoSABEZdshN7BRmTttK2c,30420
 tf_keras/src/saving/legacy/saved_model/serialized_attributes.py,sha256=nlmtIzLUBGSQU6gDKcg4-ypSRX3RbS4vPmLIhG3HSbk,15009
 tf_keras/src/saving/legacy/saved_model/utils.py,sha256=2OCwun0U8nsZvxUbv7Toq2EeC1HU32LxnLDan8cw4Dc,9953
 tf_keras/src/testing_infra/__init__.py,sha256=yrmnTOUMQ09fOgD3PD4NjpaeKz2OXCUmmoExRWhg9AY,690
@@ -585,7 +585,7 @@ tf_keras/src/utils/io_utils.py,sha256=XhCTkjwtfBc2hWSenjVdt0-2PsIc2bjJVWEP1880NU
 tf_keras/src/utils/keras_logging.py,sha256=Fv4eOMemx3Jg1hEdHIxx9GblG5YTnW1q1D1zLF3JxUE,882
 tf_keras/src/utils/kernelized_utils.py,sha256=s475SAos2zHQ1NT9AHZmbWUSahHKOhdctP6uIou0nRo,4517
 tf_keras/src/utils/kpl_test_utils.py,sha256=vnaJkySSTVhXsFEdDxNJArwXaah0yPNTK8o_3rYZvOE,7365
-tf_keras/src/utils/layer_utils.py,sha256=5SGCXE5Tc8QmCt7JHTbpVjYBkEMtdI7TQ5jC-9VLnCY,41764
+tf_keras/src/utils/layer_utils.py,sha256=cLKqiqJ2em16zZyXaXFErsL6yja28qE6kgPs2TTcdcY,43427
 tf_keras/src/utils/losses_utils.py,sha256=oPHJSNLY8U57ieQD59vnGHNavZpMpeTZtL7VIlDwwfM,16919
 tf_keras/src/utils/metrics_utils.py,sha256=feW5GoiznbQKkxmE3Url2nlXfWgvHMpJPXdGKCdiV_U,39803
 tf_keras/src/utils/mode_keys.py,sha256=_QYq58qr_b-RhvMYBYnL47NkC0G1ng8NYcVnS_IYi-A,856
@@ -606,7 +606,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.20.0.dev2025052009.dist-info/METADATA,sha256=Zg4L_5ih4xW_9SjuJ-Mqt_BVGa0Z7fg-lHfDj5_esj4,1857
-tf_keras_nightly-2.20.0.dev2025052009.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-tf_keras_nightly-2.20.0.dev2025052009.dist-info/top_level.txt,sha256=LC8FK7zHDNKxB17C6lGKvrZ_fZZGJsRiBK23SfiDegY,9
-tf_keras_nightly-2.20.0.dev2025052009.dist-info/RECORD,,
+tf_keras_nightly-2.20.0.dev2025052209.dist-info/METADATA,sha256=eHN3mvaRbwm9CXjGPFmF-LcYZHEJ35FyBUi8YkDclxM,1857
+tf_keras_nightly-2.20.0.dev2025052209.dist-info/WHEEL,sha256=zaaOINJESkSfm_4HQVc5ssNzHCPXhJm0kEUakpsEHaU,91
+tf_keras_nightly-2.20.0.dev2025052209.dist-info/top_level.txt,sha256=LC8FK7zHDNKxB17C6lGKvrZ_fZZGJsRiBK23SfiDegY,9
+tf_keras_nightly-2.20.0.dev2025052209.dist-info/RECORD,,

{tf_keras_nightly-2.20.0.dev2025052009.dist-info → tf_keras_nightly-2.20.0.dev2025052209.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any