Repository: keras-team/keras (python)

Issue:
Layernorm not supporting axis [-2, 3]
Hi, 
I wanted to normalise my output on the -2 and -3 axis, (image height and width), however, it seems that the with rms_scaling=true, the self.gamma is not broadcasted to same shape as layer input causing this error, 
```
inputs shape: (1, 1920, 1200, 3)
inv shape: (1, 1, 1, 3)
gamma_cast shape: (1920, 1200)
inv shape: (1, 1920, 1200, 3)
2024-04-30 13:50:54.238379: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: INVALID_ARGUMENT: Incompatible shapes: [1,1920,1200,3] vs. [1920,1200]
Traceback (most recent call last):
  File "C:\Users\88bbh\PycharmProjects\AI\tempt.py", line 10, in <module>
    layer(np.zeros((1, 1920, 1200, 3)))
  File "C:\Users\88bbh\PycharmProjects\AI\venv\lib\site-packages\keras\src\utils\traceback_utils.py", line 122, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "C:\Users\88bbh\PycharmProjects\AI\venv\lib\site-packages\tensorflow\python\framework\ops.py", line 5983, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.InvalidArgumentError: Exception encountered when calling LayerNormalization.call().

{{function_node __wrapped__Mul_device_/job:localhost/replica:0/task:0/device:CPU:0}} Incompatible shapes: [1,1920,1200,3] vs. [1920,1200] [Op:Mul] name: 

Arguments received by LayerNormalization.call():
  • inputs=tf.Tensor(shape=(1, 1920, 1200, 3), dtype=float32)
```

code to reproduce
```
layer = keras.layers.LayerNormalization(axis=[-3, -2], rms_scaling=True)
layer.build([None, 1920, 1200, 3])
layer(np.zeros((1, 1920, 1200, 3)))
```

the error is in layernorm call method
```
        if self.rms_scaling:
            # Calculate outputs with only variance and gamma if rms scaling
            # is enabled
            # Calculate the variance along self.axis (layer activations).
            variance = ops.var(inputs, axis=self.axis, keepdims=True)
            inv = ops.rsqrt(variance + self.epsilon)
            print("inputs shape:", inputs.shape)
            print("inv shape:", inv.shape)
            print("gamma_cast shape:", self.gamma.shape)
            print("inv shape:", (inputs * inv).shape)
            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
```

the error can be fixed by changing
```
outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
to
outputs = inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)
```
please fix it in the next update
thank you

Pull request description:
Fix bug in layernorm with rms_scaling=True
Fix bug in `LayerNormalization` with `rms_scaling` set to `True`.

Fixes #19642


Patch:
diff --git a/.github/workflows/scorecard.yml b/.github/workflows/scorecard.yml
index 1286b78b5..2fe3b5cdc 100644
--- a/.github/workflows/scorecard.yml
+++ b/.github/workflows/scorecard.yml
@@ -48,7 +48,7 @@ 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
@@ -56,6 +56,6 @@ jobs:
 
       # 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
diff --git a/keras/src/backend/jax/core.py b/keras/src/backend/jax/core.py
index 4d95bcaa7..d0b5f9c8e 100644
--- a/keras/src/backend/jax/core.py
+++ b/keras/src/backend/jax/core.py
@@ -79,8 +79,8 @@ def convert_to_numpy(x):
     if isinstance(x, jax_sparse.JAXSparse):
         x = x.todense()
     if is_tensor(x) and x.dtype == "bfloat16":
-        return np.asarray(x, ml_dtypes.bfloat16)
-    return np.asarray(x)
+        return np.array(x, dtype=ml_dtypes.bfloat16)
+    return np.array(x)
 
 
 def is_tensor(x):
diff --git a/keras/src/backend/tensorflow/core.py b/keras/src/backend/tensorflow/core.py
index 575e6dc6d..89629215f 100644
--- a/keras/src/backend/tensorflow/core.py
+++ b/keras/src/backend/tensorflow/core.py
@@ -128,7 +128,7 @@ def convert_to_numpy(x):
         x = tf.convert_to_tensor(x)
     elif isinstance(x, tf.RaggedTensor):
         x = x.to_tensor()
-    return np.asarray(x)
+    return np.array(x)
 
 
 def is_tensor(x):
diff --git a/keras/src/backend/tensorflow/numpy.py b/keras/src/backend/tensorflow/numpy.py
index 51ee5833c..74920a624 100644
--- a/keras/src/backend/tensorflow/numpy.py
+++ b/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
 
 
diff --git a/keras/src/backend/torch/linalg.py b/keras/src/backend/torch/linalg.py
index 81041782a..9a15f2478 100644
--- a/keras/src/backend/torch/linalg.py
+++ b/keras/src/backend/torch/linalg.py
@@ -8,7 +8,7 @@ from keras.src.backend.torch.core import convert_to_tensor
 
 
 def cholesky(x):
-    return torch.cholesky(x)
+    return torch.linalg.cholesky(x)
 
 
 def det(x):
diff --git a/keras/src/export/export_lib.py b/keras/src/export/export_lib.py
index 09ec14569..2ea5db1b7 100644
--- a/keras/src/export/export_lib.py
+++ b/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
 
diff --git a/keras/src/layers/layer.py b/keras/src/layers/layer.py
index 7d67339b6..3bf627873 100644
--- a/keras/src/layers/layer.py
+++ b/keras/src/layers/layer.py
@@ -1350,10 +1350,10 @@ class Layer(BackendLayer, Operation, KerasSaveable):
     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):
diff --git a/keras/src/layers/normalization/layer_normalization.py b/keras/src/layers/normalization/layer_normalization.py
index 854a07193..321118379 100644
--- a/keras/src/layers/normalization/layer_normalization.py
+++ b/keras/src/layers/normalization/layer_normalization.py
@@ -212,7 +212,9 @@ class LayerNormalization(Layer):
             variance = ops.var(inputs, axis=self.axis, keepdims=True)
             inv = ops.rsqrt(variance + self.epsilon)
 
-            outputs = inputs * inv * ops.cast(self.gamma, inputs.dtype)
+            outputs = (
+                inputs * inv * ops.cast(_broadcast(self.gamma), inputs.dtype)
+            )
         else:
             # Calculate the mean & variance along self.axis (layer activations).
             mean, variance = ops.moments(inputs, axes=self.axis, keepdims=True)
diff --git a/keras/src/layers/rnn/gru.py b/keras/src/layers/rnn/gru.py
index f489bd663..8a516d0b4 100644
--- a/keras/src/layers/rnn/gru.py
+++ b/keras/src/layers/rnn/gru.py
@@ -538,14 +538,23 @@ class GRU(RNN):
         if tree.is_nested(mask):
             mask = mask[0]
         if self.use_cudnn in ("auto", True):
-            if not self.dropout and not self.recurrent_dropout:
+            if not self.recurrent_dropout:
                 try:
+                    if self.dropout:
+                        dp_mask = self.cell.get_dropout_mask(sequences[:, 0, :])
+                        dp_mask = ops.expand_dims(dp_mask, axis=1)
+                        dp_mask = ops.broadcast_to(
+                            dp_mask, ops.shape(sequences)
+                        )
+                        dp_sequences = sequences * dp_mask
+                    else:
+                        dp_sequences = sequences
                     # Backends are allowed to specify (optionally) optimized
                     # implementation of the inner GRU loop. In the case of
                     # TF for instance, it will leverage cuDNN when feasible, and
                     # it will raise NotImplementedError otherwise.
                     out = backend.gru(
-                        sequences,
+                        dp_sequences,
                         initial_state,
                         mask,
                         kernel=self.cell.kernel,
diff --git a/keras/src/layers/rnn/lstm.py b/keras/src/layers/rnn/lstm.py
index 33055fd19..f4903655b 100644
--- a/keras/src/layers/rnn/lstm.py
+++ b/keras/src/layers/rnn/lstm.py
@@ -518,14 +518,24 @@ class LSTM(RNN):
             mask = mask[0]
 
         if self.use_cudnn in ("auto", True):
-            if not self.dropout and not self.recurrent_dropout:
+            if not self.recurrent_dropout:
                 try:
+                    if self.dropout:
+                        dp_mask = self.cell.get_dropout_mask(sequences[:, 0, :])
+                        dp_mask = ops.expand_dims(dp_mask, axis=1)
+                        dp_mask = ops.broadcast_to(
+                            dp_mask, ops.shape(sequences)
+                        )
+                        dp_sequences = sequences * dp_mask
+                    else:
+                        dp_sequences = sequences
+
                     # Backends are allowed to specify (optionally) optimized
                     # implementation of the inner LSTM loop. In the case of
                     # TF for instance, it will leverage cuDNN when feasible, and
                     # it will raise NotImplementedError otherwise.
                     out = backend.lstm(
-                        sequences,
+                        dp_sequences,
                         initial_state[0],
                         initial_state[1],
                         mask,
diff --git a/keras/src/ops/numpy.py b/keras/src/ops/numpy.py
index c11899720..9fb631473 100644
--- a/keras/src/ops/numpy.py
+++ b/keras/src/ops/numpy.py
@@ -4859,28 +4859,9 @@ class TakeAlongAxis(Operation):
         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)
 
 
diff --git a/keras/src/ops/operation_utils.py b/keras/src/ops/operation_utils.py
index ac0961da4..f5ca1857c 100644
--- a/keras/src/ops/operation_utils.py
+++ b/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:
diff --git a/keras/src/trainers/trainer.py b/keras/src/trainers/trainer.py
index f9bb1a738..af5e3a3f7 100644
--- a/keras/src/trainers/trainer.py
+++ b/keras/src/trainers/trainer.py
@@ -130,7 +130,8 @@ class Trainer:
                 wrapped in a `LossScaleOptimizer`, which will dynamically
                 scale the loss to prevent underflow.
         """
-        self.optimizer = optimizers.get(optimizer)
+        optimizer = optimizers.get(optimizer)
+        self.optimizer = optimizer
         if (
             auto_scale_loss
             and self.dtype_policy.name == "mixed_float16"
diff --git a/keras/src/utils/model_visualization.py b/keras/src/utils/model_visualization.py
index 2784a83e3..1fd180339 100644
--- a/keras/src/utils/model_visualization.py
+++ b/keras/src/utils/model_visualization.py
@@ -112,18 +112,32 @@ def make_layer_label(layer, **kwargs):
             output_shape = tree.map_structure(lambda x: x.shape, layer.output)
         except (ValueError, AttributeError):
             pass
+
+        def format_shape(shape):
+            if shape is not None:
+                if isinstance(shape, dict):
+                    shape_str = ", ".join(
+                        [f"{k}: {v}" for k, v in shape.items()]
+                    )
+                else:
+                    shape_str = f"{shape}"
+                shape_str = shape_str.replace("}", "").replace("{", "")
+            else:
+                shape_str = "?"
+            return shape_str
+
         if class_name != "InputLayer":
             cols.append(
                 (
                     '<td bgcolor="white"><font point-size="14">'
-                    f'Input shape: <b>{input_shape or "?"}</b>'
+                    f"Input shape: <b>{format_shape(input_shape)}</b>"
                     "</font></td>"
                 )
             )
         cols.append(
             (
                 '<td bgcolor="white"><font point-size="14">'
-                f'Output shape: <b>{output_shape or "?"}</b>'
+                f"Output shape: <b>{format_shape(output_shape)}</b>"
                 "</font></td>"
             )
         )
diff --git a/pip_build.py b/pip_build.py
index 6570ab8eb..66e7578ee 100644
--- a/pip_build.py
+++ b/pip_build.py
@@ -84,7 +84,6 @@ def build(root_path, is_nightly=False, rc_index=None):
     try:
         copy_source_to_build_directory(root_path)
         move_tf_keras_directory()
-        print(os.getcwd())
 
         from keras.src.version import __version__  # noqa: E402
 
diff --git a/requirements-common.txt b/requirements-common.txt
index f645c7ba9..c972212c6 100644
--- a/requirements-common.txt
+++ b/requirements-common.txt
@@ -17,3 +17,4 @@ rich
 build
 optree
 pytest-cov
+packaging
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index c759c9d18..6de2a3c25 100644
--- a/requirements.txt
+++ b/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)
diff --git a/setup.py b/setup.py
index a78f07dda..822eccfa7 100644
--- a/setup.py
+++ b/setup.py
@@ -43,6 +43,7 @@ setup(
         "h5py",
         "optree",
         "ml-dtypes",
+        "packaging",
     ],
     # Supported Python versions
     python_requires=">=3.9",


Test patch:
diff --git a/keras/src/export/export_lib_test.py b/keras/src/export/export_lib_test.py
index dd23857df..f754525f3 100644
--- a/keras/src/export/export_lib_test.py
+++ b/keras/src/export/export_lib_test.py
@@ -95,6 +95,56 @@ class ExportArchiveTest(testing.TestCase, parameterized.TestCase):
         # 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"])
     )
diff --git a/keras/src/ops/core_test.py b/keras/src/ops/core_test.py
index 29526c20a..3993b8163 100644
--- a/keras/src/ops/core_test.py
+++ b/keras/src/ops/core_test.py
@@ -633,6 +633,13 @@ class CoreOpsDtypeTest(testing.TestCase, parameterized.TestCase):
                 expected_dtype,
             )
 
+    def test_convert_to_numpy(self):
+        x = ops.array([1, 2, 3], dtype="float32")
+        y = ops.convert_to_numpy(x)
+        self.assertIsInstance(y, np.ndarray)
+        # Test assignment -- should not fail.
+        y[0] = 1.0
+
 
 class CoreOpsCallsTests(testing.TestCase):
     def test_scatter_basic_call(self):
diff --git a/keras/src/saving/saving_lib_test.py b/keras/src/saving/saving_lib_test.py
index 23ac2a52e..80f53608f 100644
--- a/keras/src/saving/saving_lib_test.py
+++ b/keras/src/saving/saving_lib_test.py
@@ -359,7 +359,9 @@ class SavingTest(testing.TestCase):
         )
         self.assertEqual(
             config_dict["compile_config"]["optimizer"],
-            "adam",
+            keras.src.saving.serialize_keras_object(
+                keras.src.optimizers.get("adam")
+            ),
         )
         self.assertEqual(
             config_dict["compile_config"]["loss"]["config"],