tf-keras-nightly 2.17.0.dev2024031909__py3-none-any.whl → 2.19.0.dev2025011410__py3-none-any.whl

tf_keras/src/mixed_precision/test_util.py

@@ -171,14 +171,14 @@ class MultiplyLayer(AssertTypeLayer):
             activity_regularizer=self._activity_regularizer, **kwargs
         )
 
-    def build(self, _):
+    def build(self, input_shape):
         self.v = self.add_weight(
             self._var_name,
             (),
             initializer="ones",
             regularizer=self._regularizer,
         )
-        self.built = True
+        super().build(input_shape)
 
     def call(self, inputs):
         self.assert_input_types(inputs)
@@ -205,7 +205,7 @@ class MultiplyLayer(AssertTypeLayer):
 class MultiplyLayerWithoutAutoCast(MultiplyLayer):
     """Same as MultiplyLayer, but does not use AutoCastVariables."""
 
-    def build(self, _):
+    def build(self, input_shape):
         dtype = self.dtype
         if dtype in ("float16", "bfloat16"):
             dtype = "float32"
@@ -214,10 +214,11 @@ class MultiplyLayerWithoutAutoCast(MultiplyLayer):
             (),
             initializer="ones",
             dtype=dtype,
-            experimental_autocast=False,
+            autocast=False,
             regularizer=self._regularizer,
         )
-        self.built = True
+        # Call Layer.build() to skip MultiplyLayer.build() which we override.
+        base_layer.Layer.build(self, input_shape)
 
     def call(self, inputs):
         self.assert_input_types(inputs)

tf_keras/src/optimizers/legacy/optimizer_v2.py

@@ -1033,6 +1033,13 @@ class OptimizerV2(tf.__internal__.tracking.Trackable):
         slot_dict = self._slots.setdefault(var_key, {})
         weight = slot_dict.get(slot_name, None)
         if weight is None:
+            # Under a mixed precision policy, variables report their "cast"
+            # dtype. However, we want to use the original dtype for slots.
+            if hasattr(var, "true_dtype"):
+                dtype = var.true_dtype
+            else:
+                dtype = var.dtype
+
             if isinstance(initializer, str) or callable(initializer):
                 initializer = initializers.get(initializer)
                 if isinstance(
@@ -1043,7 +1050,7 @@ class OptimizerV2(tf.__internal__.tracking.Trackable):
                 else:
                     slot_shape = var.shape
                 initial_value = functools.partial(
-                    initializer, shape=slot_shape, dtype=var.dtype
+                    initializer, shape=slot_shape, dtype=dtype
                 )
             else:
                 initial_value = initializer
@@ -1064,7 +1071,7 @@ class OptimizerV2(tf.__internal__.tracking.Trackable):
                 with strategy.extended.colocate_vars_with(var):
                     weight = tf.Variable(
                         name=f"{var._shared_name}/{slot_name}",
-                        dtype=var.dtype,
+                        dtype=dtype,
                         trainable=False,
                         initial_value=initial_value,
                     )

tf_keras/src/optimizers/optimizer.py

@@ -498,26 +498,28 @@ class _BaseOptimizer(tf.__internal__.tracking.AutoTrackable):
         Returns:
           An optimizer variable.
         """
+        # Under a mixed precision policy, variables report their "cast"
+        # dtype. However, we want to use the original dtype for slots.
+        if hasattr(model_variable, "true_dtype"):
+            dtype = model_variable.true_dtype
+        else:
+            dtype = model_variable.dtype
         if initial_value is None:
             if shape is None:
                 if model_variable.shape.rank is None:
                     # When the rank is None, we cannot get a concrete
                     # `model_variable.shape`, we use dynamic shape.
-                    initial_value = tf.zeros_like(
-                        model_variable, dtype=model_variable.dtype
-                    )
+                    initial_value = tf.zeros_like(model_variable, dtype=dtype)
                 else:
                     # We cannot always use `zeros_like`, because some cases
                     # the shape exists while values don't.
-                    initial_value = tf.zeros(
-                        model_variable.shape, dtype=model_variable.dtype
-                    )
+                    initial_value = tf.zeros(model_variable.shape, dtype=dtype)
             else:
-                initial_value = tf.zeros(shape, dtype=model_variable.dtype)
+                initial_value = tf.zeros(shape, dtype=dtype)
         variable = tf.Variable(
             initial_value=initial_value,
             name=f"{variable_name}/{model_variable._shared_name}",
-            dtype=model_variable.dtype,
+            dtype=dtype,
             trainable=False,
         )
         # If model_variable is a shard of a ShardedVariable, we should add a
@@ -1188,10 +1190,17 @@ class Optimizer(_BaseOptimizer):
                         self._mesh, rank=initial_value.shape.rank
                     ),
                 )
+            # Under a mixed precision policy, variables report their "cast"
+            # dtype. However, we want to use the original dtype for optimizer
+            # variables.
+            if hasattr(model_variable, "true_dtype"):
+                dtype = model_variable.true_dtype
+            else:
+                dtype = model_variable.dtype
             variable = tf.experimental.dtensor.DVariable(
                 initial_value=initial_value,
                 name=f"{variable_name}/{model_variable._shared_name}",
-                dtype=model_variable.dtype,
+                dtype=dtype,
                 trainable=False,
             )
             self._variables.append(variable)

tf_keras/src/premade_models/linear.py

@@ -156,7 +156,8 @@ class LinearModel(training.Model):
             )
         else:
             self.bias = None
-        self.built = True
+        # Call Layer.build() to skip Model.build() which we override here.
+        base_layer.Layer.build(self, input_shape)
 
     def call(self, inputs):
         result = None

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

@@ -210,7 +210,7 @@ def get_json_type(obj):
             return {
                 "class_name": "TypeSpec",
                 "type_spec": type_spec_name,
-                "serialized": obj._serialize(),
+                "serialized": _encode_tuple(obj._serialize()),
             }
         except ValueError:
             raise ValueError(

tf_keras/src/saving/saving_api.py

@@ -15,6 +15,7 @@
 """Public API surface for saving APIs."""
 
 import os
+import tempfile
 import warnings
 import zipfile
 
@@ -33,17 +34,76 @@ except ImportError:
 is_oss = True
 
 
-def _support_gcs_uri(filepath, save_format, is_oss):
-    """Supports GCS URIs through bigstore via a temporary file."""
-    gs_filepath = None
-    if str(filepath).startswith("gs://") and save_format != "tf":
-        gs_filepath = filepath
-        if not is_oss:
-            gs_filepath = filepath.replace("gs://", "/bigstore/")
-        filepath = os.path.join(
-            saving_lib.get_temp_dir(), os.path.basename(gs_filepath)
-        )
-    return gs_filepath, filepath
+class SupportReadFromRemote:
+    """Supports GCS URIs and other remote paths via a temporary file.
+
+    This is used for `.keras` and H5 files on GCS, CNS and CFS. TensorFlow
+    supports remoted saved model out of the box.
+    """
+
+    def __init__(self, filepath):
+        save_format = get_save_format(filepath, save_format=None)
+        if (
+            saving_lib.is_remote_path(filepath)
+            and not tf.io.gfile.isdir(filepath)
+            and save_format != "tf"
+        ):
+            self.temp_directory = tempfile.TemporaryDirectory()
+            gs_filepath = filepath
+            if not is_oss and str(filepath).startswith("gs://"):
+                gs_filepath = filepath.replace("gs://", "/bigstore/")
+            self.local_filepath = os.path.join(
+                self.temp_directory.name, os.path.basename(filepath)
+            )
+            tf.io.gfile.copy(gs_filepath, self.local_filepath, overwrite=True)
+        else:
+            self.temp_directory = None
+            self.local_filepath = filepath
+
+    def __enter__(self):
+        return self.local_filepath
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.temp_directory is not None:
+            self.temp_directory.cleanup()
+
+
+class SupportWriteToRemote:
+    """Supports GCS URIs and other remote paths via a temporary file.
+
+    This is used for `.keras` and H5 files on GCS, CNS and CFS. TensorFlow
+    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)
+        self.overwrite = overwrite
+        if saving_lib.is_remote_path(filepath) and save_format != "tf":
+            self.temp_directory = tempfile.TemporaryDirectory()
+            self.remote_filepath = filepath
+            if not is_oss and str(filepath).startswith("gs://"):
+                self.remote_filepath = self.remote_filepath.replace(
+                    "gs://", "/bigstore/"
+                )
+            self.local_filepath = os.path.join(
+                self.temp_directory.name, os.path.basename(filepath)
+            )
+        else:
+            self.temp_directory = None
+            self.remote_filepath = None
+            self.local_filepath = filepath
+
+    def __enter__(self):
+        return self.local_filepath
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.temp_directory is not None:
+            tf.io.gfile.copy(
+                self.local_filepath,
+                self.remote_filepath,
+                overwrite=self.overwrite,
+            )
+            self.temp_directory.cleanup()
 
 
 @keras_export("keras.saving.save_model", "keras.models.save_model")
@@ -131,46 +191,49 @@ 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 GCS URIs through bigstore via a temporary file
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-
-    # Deprecation warnings
-    if save_format == "h5":
-        warnings.warn(
-            "You are saving your model as an HDF5 file via `model.save()`. "
-            "This file format is considered legacy. "
-            "We recommend using instead the native TF-Keras format, "
-            "e.g. `model.save('my_model.keras')`.",
-            stacklevel=2,
-        )
+    # 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(
+                "You are saving your model as an HDF5 file via `model.save()`. "
+                "This file format is considered legacy. "
+                "We recommend using instead the native TF-Keras format, "
+                "e.g. `model.save('my_model.keras')`.",
+                stacklevel=2,
+            )
 
-    if save_format == "keras":
-        # If file exists and should not be overwritten.
-        try:
-            exists = os.path.exists(filepath)
-        except TypeError:
-            exists = False
-        if exists and not overwrite:
-            proceed = io_utils.ask_to_proceed_with_overwrite(filepath)
-            if not proceed:
-                return
-        if kwargs:
-            raise ValueError(
-                "The following argument(s) are not supported "
-                f"with the native TF-Keras format: {list(kwargs.keys())}"
+        if save_format == "keras":
+            # If file exists and should not be overwritten.
+            try:
+                exists = os.path.exists(local_filepath)
+            except TypeError:
+                exists = False
+            if exists and not overwrite:
+                proceed = io_utils.ask_to_proceed_with_overwrite(local_filepath)
+                if not proceed:
+                    return
+            if kwargs:
+                raise ValueError(
+                    "The following argument(s) are not supported "
+                    f"with the native TF-Keras format: {list(kwargs.keys())}"
+                )
+            saving_lib.save_model(model, local_filepath)
+        else:
+            # Legacy case
+            return legacy_sm_saving_lib.save_model(
+                model,
+                local_filepath,
+                overwrite=overwrite,
+                save_format=save_format,
+                **kwargs,
             )
-        saving_lib.save_model(model, filepath)
-    else:
-        # Legacy case
-        return legacy_sm_saving_lib.save_model(
-            model,
-            filepath,
-            overwrite=overwrite,
-            save_format=save_format,
-            **kwargs,
-        )
 
 
 @keras_export("keras.saving.load_model", "keras.models.load_model")
@@ -217,94 +280,69 @@ def load_model(
     It is recommended that you use layer attributes to
     access specific variables, e.g. `model.get_layer("dense_1").kernel`.
     """
-    # Supports GCS URIs by copying data to temporary file
-    save_format = get_save_format(filepath, save_format=None)
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-    if gs_filepath is not None:
-        tf.io.gfile.copy(gs_filepath, filepath, overwrite=True)
-
-    is_keras_zip = str(filepath).endswith(".keras") and zipfile.is_zipfile(
-        filepath
-    )
-
-    # Support for remote zip files
-    if (
-        saving_lib.is_remote_path(filepath)
-        and not tf.io.gfile.isdir(filepath)
-        and not is_keras_zip
-    ):
-        local_path = os.path.join(
-            saving_lib.get_temp_dir(), os.path.basename(filepath)
-        )
-
-        # Copy from remote to temporary local directory
-        tf.io.gfile.copy(filepath, local_path, overwrite=True)
-
-        # Switch filepath to local zipfile for loading model
-        if zipfile.is_zipfile(local_path):
-            filepath = local_path
-            is_keras_zip = True
-
-    if is_keras_zip:
-        if kwargs:
-            raise ValueError(
-                "The following argument(s) are not supported "
-                f"with the native TF-Keras format: {list(kwargs.keys())}"
+    # Supports remote paths via a temporary file
+    with SupportReadFromRemote(filepath) as local_filepath:
+        if str(local_filepath).endswith(".keras") and zipfile.is_zipfile(
+            local_filepath
+        ):
+            if kwargs:
+                raise ValueError(
+                    "The following argument(s) are not supported "
+                    f"with the native TF-Keras format: {list(kwargs.keys())}"
+                )
+            return saving_lib.load_model(
+                local_filepath,
+                custom_objects=custom_objects,
+                compile=compile,
+                safe_mode=safe_mode,
             )
-        return saving_lib.load_model(
-            filepath,
+
+        # Legacy case.
+        return legacy_sm_saving_lib.load_model(
+            local_filepath,
             custom_objects=custom_objects,
             compile=compile,
-            safe_mode=safe_mode,
+            **kwargs,
         )
 
-    # Legacy case.
-    return legacy_sm_saving_lib.load_model(
-        filepath, custom_objects=custom_objects, compile=compile, **kwargs
-    )
-
 
 def save_weights(model, filepath, overwrite=True, **kwargs):
-    # Supports GCS URIs through bigstore via a temporary file
-    save_format = get_save_format(filepath, save_format=None)
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-
-    if str(filepath).endswith(".weights.h5"):
-        # If file exists and should not be overwritten.
-        try:
-            exists = os.path.exists(filepath)
-        except TypeError:
-            exists = False
-        if exists and not overwrite:
-            proceed = io_utils.ask_to_proceed_with_overwrite(filepath)
-            if not proceed:
-                return
-        saving_lib.save_weights_only(model, filepath)
-    else:
-        legacy_sm_saving_lib.save_weights(
-            model, filepath, overwrite=overwrite, **kwargs
-        )
+    # Supports remote paths via a temporary file
+    with SupportWriteToRemote(filepath, overwrite=overwrite) as local_filepath:
+        if str(local_filepath).endswith(".weights.h5"):
+            # If file exists and should not be overwritten.
+            try:
+                exists = os.path.exists(local_filepath)
+            except TypeError:
+                exists = False
+            if exists and not overwrite:
+                proceed = io_utils.ask_to_proceed_with_overwrite(local_filepath)
+                if not proceed:
+                    return
+            saving_lib.save_weights_only(model, local_filepath)
+        else:
+            legacy_sm_saving_lib.save_weights(
+                model, local_filepath, overwrite=overwrite, **kwargs
+            )
 
 
 def load_weights(model, filepath, skip_mismatch=False, **kwargs):
-    # Supports GCS URIs by copying data to temporary file
-    save_format = get_save_format(filepath, save_format=None)
-    gs_filepath, filepath = _support_gcs_uri(filepath, save_format, is_oss)
-    if gs_filepath is not None:
-        tf.io.gfile.copy(gs_filepath, filepath, overwrite=True)
-
-    if str(filepath).endswith(".keras") and zipfile.is_zipfile(filepath):
-        saving_lib.load_weights_only(
-            model, filepath, skip_mismatch=skip_mismatch
-        )
-    elif str(filepath).endswith(".weights.h5"):
-        saving_lib.load_weights_only(
-            model, filepath, skip_mismatch=skip_mismatch
-        )
-    else:
-        return legacy_sm_saving_lib.load_weights(
-            model, filepath, skip_mismatch=skip_mismatch, **kwargs
-        )
+    # Supports remote paths via a temporary file
+    with SupportReadFromRemote(filepath) as local_filepath:
+        if str(local_filepath).endswith(".keras") and zipfile.is_zipfile(
+            local_filepath
+        ):
+            saving_lib.load_weights_only(
+                model, local_filepath, skip_mismatch=skip_mismatch
+            )
+        elif str(local_filepath).endswith(".weights.h5"):
+            saving_lib.load_weights_only(
+                model, local_filepath, skip_mismatch=skip_mismatch
+            )
+        else:
+            return legacy_sm_saving_lib.load_weights(
+                model, local_filepath, skip_mismatch=skip_mismatch, **kwargs
+            )
 
 
 def get_save_format(filepath, save_format):

tf_keras/src/saving/saving_lib.py

@@ -156,14 +156,8 @@ def save_model(model, filepath, weights_format="h5"):
             "date_saved": datetime.datetime.now().strftime("%Y-%m-%d@%H:%M:%S"),
         }
     )
-    # TODO(rameshsampath): Need a better logic for local vs remote path
-    if is_remote_path(filepath):
-        # Remote path. Zip to local memory byte io and copy to remote
-        zip_filepath = io.BytesIO()
-    else:
-        zip_filepath = filepath
     try:
-        with zipfile.ZipFile(zip_filepath, "w") as zf:
+        with zipfile.ZipFile(filepath, "w") as zf:
             with zf.open(_METADATA_FILENAME, "w") as f:
                 f.write(metadata_json.encode())
             with zf.open(_CONFIG_FILENAME, "w") as f:
@@ -195,10 +189,6 @@ def save_model(model, filepath, weights_format="h5"):
             )
             weights_store.close()
             asset_store.close()
-
-        if is_remote_path(filepath):
-            with tf.io.gfile.GFile(filepath, "wb") as f:
-                f.write(zip_filepath.getvalue())
     except Exception as e:
         raise e
     finally:

tf_keras/src/saving/serialization_lib.py

@@ -378,7 +378,7 @@ def _get_class_or_fn_config(obj):
     """Return the object's config depending on its type."""
     # Functions / lambdas:
     if isinstance(obj, types.FunctionType):
-        return obj.__name__
+        return object_registration.get_registered_name(obj)
     # All classes:
     if hasattr(obj, "get_config"):
         config = obj.get_config()
@@ -789,15 +789,6 @@ def _retrieve_class_or_fn(
                 if obj is not None:
                     return obj
 
-            # Retrieval of registered custom function in a package
-            filtered_dict = {
-                k: v
-                for k, v in custom_objects.items()
-                if k.endswith(full_config["config"])
-            }
-            if filtered_dict:
-                return next(iter(filtered_dict.values()))
-
         # Otherwise, attempt to retrieve the class object given the `module`
         # and `class_name`. Import the module, find the class.
         try:

tf_keras/src/utils/data_utils.py

@@ -1108,7 +1108,7 @@ def pad_sequences(
         maxlen = np.max(lengths)
 
     is_dtype_str = np.issubdtype(dtype, np.str_) or np.issubdtype(
-        dtype, np.unicode_
+        dtype, np.str_
     )
     if isinstance(value, str) and dtype != object and not is_dtype_str:
         raise ValueError(

tf_keras/src/utils/steps_per_execution_tuning.py

@@ -229,7 +229,7 @@ class StepsPerExecutionTuner:
 
         if current_spe >= spe_limit:
             new_spe = current_spe
-        elif current_spe == 0:
+        elif current_spe <= 0:
             new_spe = self.init_spe
 
         self._steps_per_execution.assign(np.round(new_spe))

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

@@ -110,16 +110,24 @@ def timeseries_dataset_from_array(
     timesteps to predict the next timestep, you would use:
 
     ```python
-    input_data = data[:-10]
-    targets = data[10:]
+    data = tf.range(15)
+    sequence_length = 10
+    input_data = data[:]
+    targets = data[sequence_length:]
     dataset = tf.keras.utils.timeseries_dataset_from_array(
-        input_data, targets, sequence_length=10)
+        input_data, targets, sequence_length=sequence_length
+    )
     for batch in dataset:
       inputs, targets = batch
-      assert np.array_equal(inputs[0], data[:10])  # First sequence: steps [0-9]
+      # First sequence: steps [0-9]
+      assert np.array_equal(inputs[0], data[:sequence_length])
       # Corresponding target: step 10
-      assert np.array_equal(targets[0], data[10])
+      assert np.array_equal(targets[0], data[sequence_length])
       break
+    # To view the generated dataset
+    for batch in dataset.as_numpy_iterator():
+      input, label = batch
+      print(f"Input:{input}, target:{label}")
     ```
 
     Example 3: Temporal regression for many-to-many architectures.

{tf_keras_nightly-2.17.0.dev2024031909.dist-info → tf_keras_nightly-2.19.0.dev2025011410.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: tf_keras-nightly
-Version: 2.17.0.dev2024031909
+Version: 2.19.0.dev2025011410
 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.17.0.dev
+Requires-Dist: tf-nightly~=2.19.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.