Merge branch 'keras-team:master' into master

.github/workflows/actions.yml

@@ -54,7 +54,6 @@ jobs:
           fi
           pip install -r $REQUIREMENTS_FILE --progress-bar off --upgrade
           pip uninstall -y keras keras-nightly
-          pip install tf_keras==2.16.0 --progress-bar off --upgrade
           pip install -e "." --progress-bar off --upgrade
       - name: Test applications with pytest
         if: ${{ steps.filter.outputs.applications == 'true' }}

.github/workflows/scorecard.yml

@@ -48,14 +48,14 @@ jobs:
       # Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
       # format to the repository Actions tab.
       - name: "Upload artifact"
-        uses: actions/upload-artifact@b4b15b8c7c6ac21ea08fcf65892d2ee8f75cf882 # v4.4.3
+        uses: actions/upload-artifact@6f51ac03b9356f520e9adb1b1b7802705f340c2b # v4.5.0
         with:
           name: SARIF file
           path: results.sarif
           retention-days: 5
 
       # Upload the results to GitHub's code scanning dashboard.
       - name: "Upload to code-scanning"
-        uses: github/codeql-action/upload-sarif@f09c1c0a94de965c15400f5634aa42fac8fb8f88 # v3.27.5
+        uses: github/codeql-action/upload-sarif@48ab28a6f5dbc2a99bf1e0131198dd8f1df78169 # v3.28.0
         with:
           sarif_file: results.sarif

guides/functional_api.py

@@ -179,6 +179,7 @@
 from this file, even if the code that built the model is no longer available.
 
 This saved file includes the:
+
 - model architecture
 - model weight values (that were learned during training)
 - model training config, if any (as passed to `compile()`)

keras/api/_tf_keras/keras/export/__init__.py

@@ -4,4 +4,4 @@
 since your modifications would be overwritten.
 """
 
-from keras.src.export.export_lib import ExportArchive
+from keras.src.export.saved_model import ExportArchive

keras/api/_tf_keras/keras/layers/__init__.py

@@ -4,7 +4,7 @@
 since your modifications would be overwritten.
 """
 
-from keras.src.export.export_lib import TFSMLayer
+from keras.src.export.tfsm_layer import TFSMLayer
 from keras.src.layers import deserialize
 from keras.src.layers import serialize
 from keras.src.layers.activations.activation import Activation
@@ -155,6 +155,12 @@
 from keras.src.layers.preprocessing.image_preprocessing.random_brightness import (
     RandomBrightness,
 )
+from keras.src.layers.preprocessing.image_preprocessing.random_color_degeneration import (
+    RandomColorDegeneration,
+)
+from keras.src.layers.preprocessing.image_preprocessing.random_color_jitter import (
+    RandomColorJitter,
+)
 from keras.src.layers.preprocessing.image_preprocessing.random_contrast import (
     RandomContrast,
 )
@@ -170,12 +176,21 @@
 from keras.src.layers.preprocessing.image_preprocessing.random_hue import (
     RandomHue,
 )
+from keras.src.layers.preprocessing.image_preprocessing.random_posterization import (
+    RandomPosterization,
+)
 from keras.src.layers.preprocessing.image_preprocessing.random_rotation import (
     RandomRotation,
 )
 from keras.src.layers.preprocessing.image_preprocessing.random_saturation import (
     RandomSaturation,
 )
+from keras.src.layers.preprocessing.image_preprocessing.random_sharpness import (
+    RandomSharpness,
+)
+from keras.src.layers.preprocessing.image_preprocessing.random_shear import (
+    RandomShear,
+)
 from keras.src.layers.preprocessing.image_preprocessing.random_translation import (
     RandomTranslation,
 )

keras/api/export/__init__.py

@@ -4,4 +4,4 @@
 since your modifications would be overwritten.
 """
 
-from keras.src.export.export_lib import ExportArchive
+from keras.src.export.saved_model import ExportArchive

keras/api/layers/__init__.py

@@ -4,7 +4,7 @@
 since your modifications would be overwritten.
 """
 
-from keras.src.export.export_lib import TFSMLayer
+from keras.src.export.tfsm_layer import TFSMLayer
 from keras.src.layers import deserialize
 from keras.src.layers import serialize
 from keras.src.layers.activations.activation import Activation
@@ -155,6 +155,12 @@
 from keras.src.layers.preprocessing.image_preprocessing.random_brightness import (
     RandomBrightness,
 )
+from keras.src.layers.preprocessing.image_preprocessing.random_color_degeneration import (
+    RandomColorDegeneration,
+)
+from keras.src.layers.preprocessing.image_preprocessing.random_color_jitter import (
+    RandomColorJitter,
+)
 from keras.src.layers.preprocessing.image_preprocessing.random_contrast import (
     RandomContrast,
 )
@@ -170,12 +176,21 @@
 from keras.src.layers.preprocessing.image_preprocessing.random_hue import (
     RandomHue,
 )
+from keras.src.layers.preprocessing.image_preprocessing.random_posterization import (
+    RandomPosterization,
+)
 from keras.src.layers.preprocessing.image_preprocessing.random_rotation import (
     RandomRotation,
 )
 from keras.src.layers.preprocessing.image_preprocessing.random_saturation import (
     RandomSaturation,
 )
+from keras.src.layers.preprocessing.image_preprocessing.random_sharpness import (
+    RandomSharpness,
+)
+from keras.src.layers.preprocessing.image_preprocessing.random_shear import (
+    RandomShear,
+)
 from keras.src.layers.preprocessing.image_preprocessing.random_translation import (
     RandomTranslation,
 )

keras/src/backend/common/variables.py

@@ -33,9 +33,11 @@ class Variable:
         autocast: Optional. Boolean indicating whether the variable supports
             autocasting. If `True`, the layer may first convert the variable
             to the compute data type when accessed. Defaults to `True`.
-        aggregation: Optional. String specifying how a distributed variable will
-            be aggregated. This serves as a semantic annotation, to be taken
-            into account by downstream backends or users. Defaults to `"mean"`.
+        aggregation: Optional string, one of `None`, `"none"`, `"mean"`,
+            `"sum"` or `"only_first_replica"` specifying how a distributed
+            variable will be aggregated. This serves as a semantic annotation,
+            to be taken into account by downstream backends or users. Defaults
+            to `"none"`.
         name: Optional. A unique name for the variable. Automatically generated
             if not set.
 
@@ -93,7 +95,7 @@ def __init__(
         dtype=None,
         trainable=True,
         autocast=True,
-        aggregation="mean",
+        aggregation="none",
         name=None,
     ):
         name = name or auto_name(self.__class__.__name__)
@@ -103,12 +105,21 @@ def __init__(
                 "cannot contain character `/`. "
                 f"Received: name={name}"
             )
-        if aggregation not in ("none", "mean", "sum", "only_first_replica"):
+        if aggregation not in (
+            None,
+            "none",
+            "mean",
+            "sum",
+            "only_first_replica",
+        ):
             raise ValueError(
                 "Invalid valid for argument `aggregation`. Expected "
-                "one of {'none', 'mean', 'sum', 'only_first_replica'}. "
+                "one of `None`, `'none'`, `'mean'`, `'sum'`, "
+                "`'only_first_replica'`. "
                 f"Received: aggregation={aggregation}"
             )
+        if aggregation is None:
+            aggregation = "none"
         self._name = name
         parent_path = current_path()
         if parent_path:

keras/src/backend/tensorflow/distribute_test.py

@@ -130,8 +130,8 @@ def test_variable_aggregation(self):
         with strategy.scope():
             x = np.random.random((4, 4))
             v1 = backend.Variable(x, dtype="float32")
-            self.assertEqual(v1.aggregation, "mean")
-            self.assertEqual(v1.value.aggregation, tf.VariableAggregation.MEAN)
+            self.assertEqual(v1.aggregation, "none")
+            self.assertEqual(v1.value.aggregation, tf.VariableAggregation.NONE)
 
             v2 = backend.Variable(x, dtype="float32", aggregation="sum")
             self.assertEqual(v2.aggregation, "sum")

keras/src/backend/torch/export.py

@@ -1,35 +1,128 @@
-from keras.src import layers
+import copy
+import warnings
+
+import torch
+
 from keras.src import tree
+from keras.src.export.export_utils import convert_spec_to_tensor
+from keras.src.utils.module_utils import tensorflow as tf
+from keras.src.utils.module_utils import torch_xla
 
 
 class TorchExportArchive:
     def track(self, resource):
-        if not isinstance(resource, layers.Layer):
-            raise ValueError(
-                "Invalid resource type. Expected an instance of a "
-                "JAX-based Keras `Layer` or `Model`. "
-                f"Received instead an object of type '{type(resource)}'. "
-                f"Object received: {resource}"
-            )
+        raise NotImplementedError(
+            "`track` is not implemented in the torch backend. Use"
+            "`track_and_add_endpoint` instead."
+        )
 
-        if isinstance(resource, layers.Layer):
-            # Variables in the lists below are actually part of the trackables
-            # that get saved, because the lists are created in __init__.
-            variables = resource.variables
-            trainable_variables = resource.trainable_variables
-            non_trainable_variables = resource.non_trainable_variables
-            self._tf_trackable.variables += tree.map_structure(
-                self._convert_to_tf_variable, variables
-            )
-            self._tf_trackable.trainable_variables += tree.map_structure(
-                self._convert_to_tf_variable, trainable_variables
-            )
-            self._tf_trackable.non_trainable_variables += tree.map_structure(
-                self._convert_to_tf_variable, non_trainable_variables
+    def add_endpoint(self, name, fn, input_signature, **kwargs):
+        raise NotImplementedError(
+            "`add_endpoint` is not implemented in the torch backend. Use"
+            "`track_and_add_endpoint` instead."
+        )
+
+    def track_and_add_endpoint(self, name, resource, input_signature, **kwargs):
+        # Disable false alarms related to lifting parameters.
+        warnings.filterwarnings("ignore", message=".*created when tracing.*")
+        warnings.filterwarnings(
+            "ignore", message=".*Unable to find the path of the module.*"
+        )
+
+        if not isinstance(resource, torch.nn.Module):
+            raise TypeError(
+                "`resource` must be an instance of `torch.nn.Module`. "
+                f"Received: resource={resource} (of type {type(resource)})"
             )
 
-    def add_endpoint(self, name, fn, input_signature=None, **kwargs):
-        # TODO: torch-xla?
-        raise NotImplementedError(
-            "`add_endpoint` is not implemented in the torch backend."
+        sample_inputs = tree.map_structure(
+            lambda x: convert_spec_to_tensor(x, replace_none_number=1),
+            input_signature,
+        )
+        sample_inputs = tuple(sample_inputs)
+
+        # Ref: torch_xla.tf_saved_model_integration
+        # TODO: Utilize `dynamic_shapes`
+        exported = torch.export.export(
+            resource, sample_inputs, dynamic_shapes=None, strict=False
+        )
+        options = torch_xla.stablehlo.StableHLOExportOptions(
+            override_tracing_arguments=sample_inputs
+        )
+        stablehlo_model = torch_xla.stablehlo.exported_program_to_stablehlo(
+            exported, options
+        )
+        state_dict_keys = list(stablehlo_model._bundle.state_dict.keys())
+
+        # Remove unused variables.
+        for k in state_dict_keys:
+            if "lifted" not in k:
+                stablehlo_model._bundle.state_dict.pop(k)
+
+        bundle = copy.deepcopy(stablehlo_model._bundle)
+        bundle.state_dict = {
+            k: tf.Variable(v, trainable=False, name=k)
+            for k, v in bundle.state_dict.items()
+        }
+        bundle.additional_constants = [
+            tf.Variable(v, trainable=False) for v in bundle.additional_constants
+        ]
+
+        # Track variables in `bundle` for `write_out`.
+        self._tf_trackable.variables += (
+            list(bundle.state_dict.values()) + bundle.additional_constants
+        )
+
+        # Ref: torch_xla.tf_saved_model_integration.save_stablehlo_graph_as_tf
+        def make_tf_function(func, bundle):
+            from tensorflow.compiler.tf2xla.python import xla as tfxla
+
+            def _get_shape_with_dynamic(signature):
+                shape = copy.copy(signature.shape)
+                for i in signature.dynamic_dims:
+                    shape[i] = None
+                return shape
+
+            def _extract_call_parameters(args, meta, bundle):
+                call_args = []
+                if meta.input_pytree_spec is not None:
+                    args = tree.flatten(args)
+                for loc in meta.input_locations:
+                    if loc.type_ == torch_xla.stablehlo.VariableType.PARAMETER:
+                        call_args.append(bundle.state_dict[loc.name])
+                    elif loc.type_ == torch_xla.stablehlo.VariableType.CONSTANT:
+                        call_args.append(
+                            bundle.additional_constants[loc.position]
+                        )
+                    else:
+                        call_args.append(args[loc.position])
+                return call_args
+
+            def inner(*args):
+                Touts = [sig.dtype for sig in func.meta.output_signature]
+                Souts = [
+                    _get_shape_with_dynamic(sig)
+                    for sig in func.meta.output_signature
+                ]
+                call_args = _extract_call_parameters(args, func.meta, bundle)
+                results = tfxla.call_module(
+                    tuple(call_args),
+                    version=5,
+                    Tout=Touts,  # dtype information
+                    Sout=Souts,  # Shape information
+                    function_list=[],
+                    module=func.bytecode,
+                )
+                if len(Souts) == 1:
+                    results = results[0]
+                return results
+
+            return inner
+
+        decorated_fn = tf.function(
+            make_tf_function(
+                stablehlo_model._bundle.stablehlo_funcs[0], bundle
+            ),
+            input_signature=input_signature,
         )
+        return decorated_fn

keras/src/callbacks/backup_and_restore.py

@@ -37,6 +37,7 @@ class BackupAndRestore(Callback):
     >>> callback = keras.callbacks.BackupAndRestore(backup_dir="/tmp/backup")
     >>> model = keras.models.Sequential([keras.layers.Dense(10)])
     >>> model.compile(keras.optimizers.SGD(), loss='mse')
+    >>> model.build(input_shape=(None, 20))
     >>> try:
     ...   model.fit(np.arange(100).reshape(5, 20), np.zeros(5), epochs=10,
     ...             batch_size=1, callbacks=[callback, InterruptingCallback()],

keras/src/export/__init__.py

@@ -1 +1,4 @@
-from keras.src.export.export_lib import ExportArchive
+from keras.src.export.onnx import export_onnx
+from keras.src.export.saved_model import ExportArchive
+from keras.src.export.saved_model import export_saved_model
+from keras.src.export.tfsm_layer import TFSMLayer