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

tf_keras/src/metrics/confusion_metrics.py

@@ -1471,9 +1471,10 @@ class AUC(base_metric.Metric):
                 # label_weights should be of length equal to the number of
                 # labels.
                 shapes.append((self.label_weights, ("L",)))
-                tf.debugging.assert_shapes(
-                    shapes, message="Number of labels is not consistent."
-                )
+
+            tf.debugging.assert_shapes(
+                shapes, message="Number of labels is not consistent."
+            )
 
         # Only forward label_weights to update_confusion_matrix_variables when
         # multi_label is False. Otherwise the averaging of individual label AUCs
@@ -1611,13 +1612,59 @@ class AUC(base_metric.Metric):
             )
             x = fp_rate
             y = recall
-        else:  # curve == 'PR'.
+        elif self.curve == metrics_utils.AUCCurve.PR:
             precision = tf.math.divide_no_nan(
                 self.true_positives,
                 tf.math.add(self.true_positives, self.false_positives),
             )
             x = recall
             y = precision
+        else:  # curve == 'PR_GAIN'.
+            # Due to the hyperbolic transform, this formula is less robust than
+            # ROC or PR values. In particular
+            # 1) Both measures diverge when there are no negative examples;
+            # 2) Both measures diverge when there are no true positives;
+            # 3) Recall gain becomes negative when the recall is lower than the
+            # label average (i.e. when more negative examples are classified
+            # positive than real positives).
+            #
+            # We ignore case 1 as it is easily communicated. For case 2 we set
+            # recall_gain to 0 and precision_gain to 1. For case 3 we set the
+            # recall_gain to 0. These fixes will result in an overastimation of
+            # the AUC for estimateors that are anti-correlated with the label
+            # (at some thresholds).
+            #
+            # The scaling factor $\frac{P}{N}$ that is used to form both
+            # gain values.
+            scaling_factor = tf.math.divide_no_nan(
+                tf.math.add(self.true_positives, self.false_negatives),
+                tf.math.add(self.false_positives, self.true_negatives),
+            )
+
+            recall_gain = 1.0 - scaling_factor * tf.math.divide_no_nan(
+                self.false_negatives, self.true_positives
+            )
+            precision_gain = 1.0 - scaling_factor * tf.math.divide_no_nan(
+                self.false_positives, self.true_positives
+            )
+            # Handle case 2.
+            recall_gain = tf.where(
+                tf.equal(self.true_positives, 0.0),
+                tf.zeros_like(recall_gain),
+                recall_gain,
+            )
+            precision_gain = tf.where(
+                tf.equal(self.true_positives, 0.0),
+                tf.ones_like(precision_gain),
+                precision_gain,
+            )
+            # Handle case 3.
+            recall_gain = tf.math.maximum(
+                recall_gain, tf.zeros_like(recall_gain)
+            )
+
+            x = recall_gain
+            y = precision_gain
 
         # Find the rectangle heights based on `summation_method`.
         if (

tf_keras/src/models/sharpness_aware_minimization.py

@@ -72,17 +72,27 @@ class SharpnessAwareMinimization(Model):
         if self.num_batch_splits is not None:
             x_split = tf.split(x, self.num_batch_splits)
             y_split = tf.split(y, self.num_batch_splits)
+            # Split the sample weight if it is provided.
+            if sample_weight is not None:
+                sample_weight_split = tf.split(
+                    sample_weight, self.num_batch_splits
+                )
+            else:
+                sample_weight_split = [None] * self.num_batch_splits
         else:
             x_split = [x]
             y_split = [y]
+            sample_weight_split = [sample_weight]
 
         gradients_all_batches = []
         pred_all_batches = []
-        for x_batch, y_batch in zip(x_split, y_split):
+        for x_batch, y_batch, sample_weight_batch in zip(
+            x_split, y_split, sample_weight_split
+        ):
             epsilon_w_cache = []
             with tf.GradientTape() as tape:
-                pred = self.model(x_batch)
-                loss = self.compiled_loss(y_batch, pred)
+                pred = self(x_batch)
+                loss = self.compiled_loss(y_batch, pred, sample_weight_batch)
             pred_all_batches.append(pred)
             trainable_variables = self.model.trainable_variables
             gradients = tape.gradient(loss, trainable_variables)
@@ -98,8 +108,8 @@ class SharpnessAwareMinimization(Model):
                 epsilon_w_cache.append(epsilon_w)
 
             with tf.GradientTape() as tape:
-                pred = self(x_batch)
-                loss = self.compiled_loss(y_batch, pred)
+                pred = self(x_batch, training=True)
+                loss = self.compiled_loss(y_batch, pred, sample_weight_batch)
             gradients = tape.gradient(loss, trainable_variables)
             if len(gradients_all_batches) == 0:
                 for gradient in gradients:
@@ -127,7 +137,7 @@ class SharpnessAwareMinimization(Model):
         self.compiled_metrics.update_state(y, pred, sample_weight)
         return {m.name: m.result() for m in self.metrics}
 
-    def call(self, inputs):
+    def call(self, inputs, **kwargs):
         """Forward pass of SAM.
 
         SAM delegates the forward pass call to the wrapped model.
@@ -138,7 +148,7 @@ class SharpnessAwareMinimization(Model):
         Returns:
           A Tensor, the outputs of the wrapped model for given `inputs`.
         """
-        return self.model(inputs)
+        return self.model(inputs, **kwargs)
 
     def get_config(self):
         config = super().get_config()

tf_keras/src/preprocessing/sequence.py

@@ -365,7 +365,7 @@ def skipgrams(
         random.shuffle(words)
 
         couples += [
-            [words[i % len(words)], random.randint(1, vocabulary_size - 1)]
+            [words[i % len(words)], random.randint(1, int(vocabulary_size - 1))]
             for i in range(num_negative_samples)
         ]
         if categorical:
@@ -375,7 +375,7 @@ def skipgrams(
 
     if shuffle:
         if seed is None:
-            seed = random.randint(0, 10e6)
+            seed = random.randint(0, int(10e6))
         random.seed(seed)
         random.shuffle(couples)
         random.seed(seed)

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/saving/saving_api.py

@@ -24,6 +24,7 @@ from tensorflow.python.util.tf_export import keras_export
 
 from tf_keras.src.saving import saving_lib
 from tf_keras.src.saving.legacy import save as legacy_sm_saving_lib
+from tf_keras.src.saving.legacy import saving_utils
 from tf_keras.src.utils import io_utils
 
 try:
@@ -75,8 +76,7 @@ class SupportWriteToRemote:
     supports remoted saved model out of the box.
     """
 
-    def __init__(self, filepath, overwrite=True, save_format=None):
-        save_format = get_save_format(filepath, save_format=save_format)
+    def __init__(self, filepath, overwrite, save_format):
         self.overwrite = overwrite
         if saving_lib.is_remote_path(filepath) and save_format != "tf":
             self.temp_directory = tempfile.TemporaryDirectory()
@@ -191,14 +191,14 @@ def save_model(model, filepath, overwrite=True, save_format=None, **kwargs):
     when loading the model. See the `custom_objects` argument in
     `tf.keras.saving.load_model`.
     """
+    save_format = get_save_format(filepath, save_format)
+
     # Supports remote paths via a temporary file
     with SupportWriteToRemote(
         filepath,
         overwrite=overwrite,
         save_format=save_format,
     ) as local_filepath:
-        save_format = get_save_format(filepath, save_format)
-
         # Deprecation warnings
         if save_format == "h5":
             warnings.warn(
@@ -307,8 +307,12 @@ def load_model(
 
 
 def save_weights(model, filepath, overwrite=True, **kwargs):
+    save_format = get_save_weights_format(filepath)
+
     # Supports remote paths via a temporary file
-    with SupportWriteToRemote(filepath, overwrite=overwrite) as local_filepath:
+    with SupportWriteToRemote(
+        filepath, overwrite=overwrite, save_format=save_format
+    ) as local_filepath:
         if str(local_filepath).endswith(".weights.h5"):
             # If file exists and should not be overwritten.
             try:
@@ -386,3 +390,12 @@ def get_save_format(filepath, save_format):
     else:
         return "h5"
 
+
+def get_save_weights_format(filepath):
+    filepath = io_utils.path_to_string(filepath)
+    filepath_is_h5 = saving_utils.is_hdf5_filepath(filepath)
+    if filepath_is_h5:
+        return "h5"
+    else:
+        return "tf"
+

tf_keras/src/saving/saving_lib.py

@@ -639,7 +639,7 @@ class NpzIOStore:
                 self.f = archive.open(root_path, mode="r")
             else:
                 self.f = open(root_path, mode="rb")
-            self.contents = np.load(self.f, allow_pickle=True)
+            self.contents = np.load(self.f, allow_pickle=False)
 
     def make(self, path):
         if not path:

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/src/utils/metrics_utils.py

@@ -237,6 +237,7 @@ class AUCCurve(Enum):
 
     ROC = "ROC"
     PR = "PR"
+    PR_GAIN = "PR_GAIN"
 
     @staticmethod
     def from_str(key):
@@ -244,10 +245,12 @@ class AUCCurve(Enum):
             return AUCCurve.PR
         elif key in ("roc", "ROC"):
             return AUCCurve.ROC
+        elif key in ("pr_gain", "prgain", "PR_GAIN", "PRGAIN"):
+            return AUCCurve.PR_GAIN
         else:
             raise ValueError(
                 f'Invalid AUC curve value: "{key}". '
-                'Expected values are ["PR", "ROC"]'
+                'Expected values are ["PR", "ROC", "PR_GAIN"]'
             )
 
 

tf_keras/src/utils/tf_utils.py

@@ -78,9 +78,9 @@ def get_random_seed():
       the random seed as an integer.
     """
     if getattr(backend._SEED_GENERATOR, "generator", None):
-        return backend._SEED_GENERATOR.generator.randint(1, 1e9)
+        return backend._SEED_GENERATOR.generator.randint(1, int(1e9))
     else:
-        return random.randint(1, 1e9)
+        return random.randint(1, int(1e9))
 
 
 def is_tensor_or_tensor_list(v):

{tf_keras_nightly-2.19.0.dev2024121210.dist-info → tf_keras_nightly-2.21.0.dev2025123010.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.4
 Name: tf_keras-nightly
-Version: 2.19.0.dev2024121210
+Version: 2.21.0.dev2025123010
 Summary: Deep learning for humans.
 Home-page: https://keras.io/
 Download-URL: https://github.com/keras-team/tf-keras/tags
@@ -26,7 +26,18 @@ Classifier: Topic :: Software Development
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
 Requires-Python: >=3.9
-Requires-Dist: tf-nightly~=2.19.0.dev
+Requires-Dist: tf-nightly~=2.21.0.dev
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: download-url
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 TF-Keras is a deep learning API written in Python,
 running on top of the machine learning platform TensorFlow.