Merge branch 'master' of github.com:keras-team/keras

Author: fchollet

.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@5d5d22a31266ced268874388b861e4b58bb5c2f3 # v4.3.1
+        uses: actions/upload-artifact@65462800fd760344b1a7b4382951275a0abb4808 # v4.3.3
         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@1b1aada464948af03b950897e5eb522f92603cc2 # v3.24.9
+        uses: github/codeql-action/upload-sarif@d39d31e687223d841ef683f52467bd88e9b21c14 # v3.25.3
         with:
           sarif_file: results.sarif

keras/src/backend/tensorflow/numpy.py

@@ -1950,6 +1950,10 @@ def take(x, indices, axis=None):
 
 
 def take_along_axis(x, indices, axis=None):
+    from keras.src.ops.operation_utils import (
+        compute_take_along_axis_output_shape,
+    )
+
     x = convert_to_tensor(x)
     indices = convert_to_tensor(indices, "int64")
     if axis is None:
@@ -1959,7 +1963,13 @@ def take_along_axis(x, indices, axis=None):
                 f"Received: indices.shape={indices.shape}"
             )
         return take_along_axis(tf.reshape(x, [-1]), indices, 0)
-    rank = tf.rank(x)
+
+    # Compute the static output shape as later on, all shapes manipulations
+    # use dynamic shapes.
+    static_output_shape = compute_take_along_axis_output_shape(
+        x.shape, indices.shape, axis
+    )
+    rank = x.ndim
     static_axis = axis
     axis = axis + rank if axis < 0 else axis
 
@@ -1981,9 +1991,6 @@ def take_along_axis(x, indices, axis=None):
     x = tf.broadcast_to(x, x_shape)
     indices = tf.broadcast_to(indices, indices_shape)
 
-    # Save indices shape so we can restore it later.
-    possible_result_shape = indices.shape
-
     # Correct the indices using "fill" mode which is the same as in jax
     indices = tf.where(indices < 0, indices + x_shape[static_axis], indices)
 
@@ -1998,7 +2005,7 @@ def take_along_axis(x, indices, axis=None):
     result = tf.gather(x, indices, batch_dims=1)
     result = tf.reshape(result, indices_shape)
     result = swapaxes(result, static_axis, -1)
-    result.set_shape(possible_result_shape)
+    result.set_shape(static_output_shape)
     return result
 
 

keras/src/export/export_lib.py

@@ -635,16 +635,31 @@ def _get_save_spec(model):
     if not shapes_dict:
         return None
 
+    def make_tensor_spec(structure):
+        # We need to turn wrapper structures like TrackingDict or _DictWrapper
+        # into plain Python structures because they don't work with jax2tf/JAX.
+        if isinstance(structure, dict):
+            return {k: make_tensor_spec(v) for k, v in structure.items()}
+        if isinstance(structure, (list, tuple)):
+            if all(isinstance(d, (int, type(None))) for d in structure):
+                return tf.TensorSpec(
+                    shape=(None,) + structure[1:], dtype=model.input_dtype
+                )
+            result = [make_tensor_spec(v) for v in structure]
+            return tuple(result) if isinstance(structure, tuple) else result
+        else:
+            raise ValueError(
+                f"Unsupported type {type(structure)} for {structure}"
+            )
+
     if len(shapes_dict) == 1:
-        shape = list(shapes_dict.values())[0]
-        shape = (None,) + shape[1:]
-        return tf.TensorSpec(shape=shape, dtype=model.input_dtype)
+        value = list(shapes_dict.values())[0]
+        return make_tensor_spec(value)
 
     specs = {}
-    for key, shape in shapes_dict.items():
+    for key, value in shapes_dict.items():
         key = key.rstrip("_shape")
-        shape = (None,) + shape[1:]
-        specs[key] = tf.TensorSpec(shape=shape, dtype=model.input_dtype)
+        specs[key] = make_tensor_spec(value)
 
     return specs
 

keras/src/export/export_lib_test.py

@@ -95,6 +95,56 @@ def call(self, inputs):
         # Test with a different batch size
         revived_model.serve(tf.random.normal((6, 10)))
 
+    @parameterized.named_parameters(
+        named_product(struct_type=["tuple", "array", "dict"])
+    )
+    def test_model_with_input_structure(self, struct_type):
+
+        class TupleModel(models.Model):
+
+            def call(self, inputs):
+                x, y = inputs
+                return ops.add(x, y)
+
+        class ArrayModel(models.Model):
+
+            def call(self, inputs):
+                x = inputs[0]
+                y = inputs[1]
+                return ops.add(x, y)
+
+        class DictModel(models.Model):
+
+            def call(self, inputs):
+                x = inputs["x"]
+                y = inputs["y"]
+                return ops.add(x, y)
+
+        if struct_type == "tuple":
+            model = TupleModel()
+            ref_input = (tf.random.normal((3, 10)), tf.random.normal((3, 10)))
+        elif struct_type == "array":
+            model = ArrayModel()
+            ref_input = [tf.random.normal((3, 10)), tf.random.normal((3, 10))]
+        elif struct_type == "dict":
+            model = DictModel()
+            ref_input = {
+                "x": tf.random.normal((3, 10)),
+                "y": tf.random.normal((3, 10)),
+            }
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "exported_model")
+        ref_output = model(tree.map_structure(ops.convert_to_tensor, ref_input))
+
+        export_lib.export_model(model, temp_filepath)
+        revived_model = tf.saved_model.load(temp_filepath)
+        self.assertAllClose(ref_output, revived_model.serve(ref_input))
+        # Test with a different batch size
+        bigger_input = tree.map_structure(
+            lambda x: tf.concat([x, x], axis=0), ref_input
+        )
+        revived_model.serve(bigger_input)
+
     @parameterized.named_parameters(
         named_product(model_type=["sequential", "functional", "subclass"])
     )

keras/src/layers/layer.py

@@ -1350,10 +1350,10 @@ def __str__(self):
     def __setattr__(self, name, value):
         # Track Variables, Layers, Metrics, SeedGenerators.
         name, value = self._setattr_hook(name, value)
-        if hasattr(self, "_tracker"):
+        if name != "_tracker":
+            if not hasattr(self, "_tracker"):
+                self._initialize_tracker()
             value = self._tracker.track(value)
-        elif name != "_tracker":
-            self._initialize_tracker()
         return super().__setattr__(name, value)
 
     def _check_super_called(self):

keras/src/ops/numpy.py

@@ -4859,28 +4859,9 @@ def call(self, x, indices):
         return backend.numpy.take_along_axis(x, indices, axis=self.axis)
 
     def compute_output_spec(self, x, indices):
-        x_shape = list(x.shape)
-        indices_shape = list(indices.shape)
-        if self.axis is None:
-            x_shape = [None] if None in x_shape else [int(np.prod(x_shape))]
-
-        if len(x_shape) != len(indices_shape):
-            raise ValueError(
-                "`x` and `indices` must have the same number of dimensions, "
-                f"but receive shape {x_shape} and {indices_shape}."
-            )
-
-        del x_shape[self.axis]
-        del indices_shape[self.axis]
-        output_shape = broadcast_shapes(x_shape, indices_shape)
-        size_on_axis = indices.shape[self.axis]
-        if self.axis == -1:
-            output_shape = output_shape + [size_on_axis]
-        elif self.axis >= 0:
-            output_shape.insert(self.axis, size_on_axis)
-        else:
-            output_shape.insert(self.axis + 1, size_on_axis)
-
+        output_shape = operation_utils.compute_take_along_axis_output_shape(
+            x.shape, indices.shape, self.axis
+        )
         return KerasTensor(output_shape, dtype=x.dtype)
 
 

keras/src/ops/operation_utils.py

@@ -349,6 +349,25 @@ def compute_transpose_output_shape(input_shape, axes):
     return tuple(input_shape[ax] for ax in axes)
 
 
+def compute_take_along_axis_output_shape(input_shape, indices_shape, axis):
+    input_shape = list(input_shape)
+    indices_shape = list(indices_shape)
+    if axis is None:
+        input_shape = (
+            [None] if None in input_shape else [int(np.prod(input_shape))]
+        )
+
+    if len(input_shape) != len(indices_shape):
+        raise ValueError(
+            "`x` and `indices` must have the same number of dimensions, "
+            f"but receive shape {input_shape} and {indices_shape}."
+        )
+
+    input_shape[axis] = indices_shape[axis]
+    output_shape = broadcast_shapes(input_shape, indices_shape)
+    return output_shape
+
+
 def reduce_shape(shape, axis=None, keepdims=False):
     shape = list(shape)
     if axis is None:

requirements-common.txt

@@ -17,3 +17,4 @@ rich
 build
 optree
 pytest-cov
+packaging

requirements.txt

@@ -1,5 +1,6 @@
 # Tensorflow.
-tensorflow-cpu~=2.16.1  # Pin to TF 2.16
+tensorflow-cpu~=2.16.1;sys_platform != 'darwin'  # Pin to TF 2.16
+tensorflow~=2.16.1;sys_platform == 'darwin'
 
 # Torch.
 # TODO: Pin to < 2.3.0 (GitHub issue #19602)

setup.py

@@ -43,6 +43,7 @@ def get_version(rel_path):
         "h5py",
         "optree",
         "ml-dtypes",
+        "packaging",
     ],
     # Supported Python versions
     python_requires=">=3.9",