Repository: keras-team/keras (python)

Issue:
Keras 3.0 load h5 model with Orthogonal initializer fails
Hi guys, 

I'm trying to load an h5 model that was working in earlier versions.

* This is a small part of the h5 file, where you can see (last part of the snippet) a recurrent initializer with a classname of **Orthogonal**.

```
{"name": "decoder_gru0", "class_name": "GRU", "config": {"name": "decoder_gru0", "trainable": true, "return_sequences": true, "return_state": false, "go_backwards": false, "stateful": false, "unroll": false, "implementation": 0, "units": 488, "activation": "tanh", "recurrent_activation": "hard_sigmoid", "use_bias": true, "kernel_initializer": {"class_name": "VarianceScaling", "config": {"scale": 1.0, "mode": "fan_avg", "distribution": "uniform", "seed": null}}, "recurrent_initializer": {"class_name": "Orthogonal", "config": {"gain": 1.0, "seed": null}}, "bias_initializer": {"class_name": "Zeros", "config": {}}
```


* The error returned is:


```bash
  File "..../keras/src/initializers/__init__.py", line 118, in get
    raise ValueError(
ValueError: Could not interpret initializer identifier: {'class_name': 'Orthogonal', 'config': {'gain': 1.0, 'seed': None}}
```

## Addition

I then added the Orthogonal initializer to the custom objects, and it seems to go further, but gets stuck here:

```bash
    raise ValueError(
ValueError: Unrecognized keyword arguments passed to GRU: {'implementation': 0}
```

Any ideas on how to fix this @mehtamansi29 ?

Pull request description:
Fix the compatibility issues of `Orthogonal` and `GRU`
Fix #19828 

Just two simple fixes. Tests have been added.

Patch:
diff --git a/keras/src/backend/jax/trainer.py b/keras/src/backend/jax/trainer.py
index dd8052519..cab833368 100644
--- a/keras/src/backend/jax/trainer.py
+++ b/keras/src/backend/jax/trainer.py
@@ -63,6 +63,7 @@ class JAXTrainer(base_trainer.Trainer):
             y=y,
             y_pred=y_pred,
             sample_weight=sample_weight,
+            training=training,
         )
         if losses:
             self._losses_override.clear()
diff --git a/keras/src/backend/numpy/trainer.py b/keras/src/backend/numpy/trainer.py
index 68eeba340..6d40982be 100644
--- a/keras/src/backend/numpy/trainer.py
+++ b/keras/src/backend/numpy/trainer.py
@@ -28,8 +28,8 @@ class NumpyTrainer(base_trainer.Trainer):
             y_pred = self(x, training=False)
         else:
             y_pred = self(x)
-        loss = self.compute_loss(
-            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
+        loss = self._compute_loss(
+            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight, training=False
         )
         self._loss_tracker.update_state(
             loss, sample_weight=tree.flatten(x)[0].shape[0]
diff --git a/keras/src/backend/tensorflow/trainer.py b/keras/src/backend/tensorflow/trainer.py
index 0f7498f92..44582760f 100644
--- a/keras/src/backend/tensorflow/trainer.py
+++ b/keras/src/backend/tensorflow/trainer.py
@@ -51,8 +51,12 @@ class TensorFlowTrainer(base_trainer.Trainer):
                 y_pred = self(x, training=True)
             else:
                 y_pred = self(x)
-            loss = self.compute_loss(
-                x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
+            loss = self._compute_loss(
+                x=x,
+                y=y,
+                y_pred=y_pred,
+                sample_weight=sample_weight,
+                training=True,
             )
             self._loss_tracker.update_state(
                 loss, sample_weight=tf.shape(tree.flatten(x)[0])[0]
@@ -78,8 +82,8 @@ class TensorFlowTrainer(base_trainer.Trainer):
             y_pred = self(x, training=False)
         else:
             y_pred = self(x)
-        loss = self.compute_loss(
-            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
+        loss = self._compute_loss(
+            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight, training=False
         )
         self._loss_tracker.update_state(
             loss, sample_weight=tf.shape(tree.flatten(x)[0])[0]
@@ -601,8 +605,8 @@ class TensorFlowTrainer(base_trainer.Trainer):
         self, y, y_pred, sample_weight=None, regularization_losses=None
     ):
         warnings.warn(
-            "`model.compiled_loss()` is deprecated. "
-            "Instead, use `model.compute_loss(x, y, y_pred, sample_weight)`.",
+            "`model.compiled_loss()` is deprecated. Instead, use "
+            "`model.compute_loss(x, y, y_pred, sample_weight, training)`.",
         )
         return self.compute_loss(
             x=None, y=y, y_pred=y_pred, sample_weight=sample_weight
@@ -610,8 +614,8 @@ class TensorFlowTrainer(base_trainer.Trainer):
 
     def loss(self, y, y_pred, sample_weight=None):
         warnings.warn(
-            "`model.loss` is deprecated. "
-            "Instead, use `model.compute_loss(x, y, y_pred, sample_weight)`.",
+            "`model.loss()` is deprecated. Instead, use "
+            "`model.compute_loss(x, y, y_pred, sample_weight, training)`.",
         )
         return self.compute_loss(
             x=None, y=y, y_pred=y_pred, sample_weight=sample_weight
diff --git a/keras/src/backend/torch/trainer.py b/keras/src/backend/torch/trainer.py
index 1572d3c90..bf0e133d6 100644
--- a/keras/src/backend/torch/trainer.py
+++ b/keras/src/backend/torch/trainer.py
@@ -49,8 +49,8 @@ class TorchTrainer(base_trainer.Trainer):
         # for the weights from the previous train step.
         self.zero_grad()
 
-        loss = self.compute_loss(
-            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
+        loss = self._compute_loss(
+            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight, training=True
         )
         self._loss_tracker.update_state(
             loss, sample_weight=tree.flatten(x)[0].shape[0]
@@ -85,8 +85,8 @@ class TorchTrainer(base_trainer.Trainer):
             y_pred = self(x, training=False)
         else:
             y_pred = self(x)
-        loss = self.compute_loss(
-            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight
+        loss = self._compute_loss(
+            x=x, y=y, y_pred=y_pred, sample_weight=sample_weight, training=False
         )
         self._loss_tracker.update_state(
             loss, sample_weight=tree.flatten(x)[0].shape[0]
diff --git a/keras/src/initializers/__init__.py b/keras/src/initializers/__init__.py
index af46b7ff9..e7cf6f76e 100644
--- a/keras/src/initializers/__init__.py
+++ b/keras/src/initializers/__init__.py
@@ -49,6 +49,7 @@ ALL_OBJECTS_DICT.update(
         "uniform": RandomUniform,
         "normal": RandomNormal,
         "orthogonal": OrthogonalInitializer,
+        "Orthogonal": OrthogonalInitializer,  # Legacy
         "one": Ones,
         "zero": Zeros,
     }
diff --git a/keras/src/layers/normalization/unit_normalization.py b/keras/src/layers/normalization/unit_normalization.py
index 1dbf97c74..76eb923b6 100644
--- a/keras/src/layers/normalization/unit_normalization.py
+++ b/keras/src/layers/normalization/unit_normalization.py
@@ -14,7 +14,7 @@ class UnitNormalization(Layer):
 
     >>> data = np.arange(6).reshape(2, 3)
     >>> normalized_data = keras.layers.UnitNormalization()(data)
-    >>> print(np.sum(normalized_data[0, :] ** 2)
+    >>> np.sum(normalized_data[0, :] ** 2)
     1.0
 
     Args:
diff --git a/keras/src/layers/rnn/gru.py b/keras/src/layers/rnn/gru.py
index 8a516d0b4..7372d769b 100644
--- a/keras/src/layers/rnn/gru.py
+++ b/keras/src/layers/rnn/gru.py
@@ -500,6 +500,7 @@ class GRU(RNN):
             trainable=kwargs.get("trainable", True),
             name="gru_cell",
             seed=seed,
+            implementation=kwargs.pop("implementation", 2),
         )
         super().__init__(
             cell,
diff --git a/keras/src/ops/core.py b/keras/src/ops/core.py
index 5d7fdf2ff..db91f1e78 100644
--- a/keras/src/ops/core.py
+++ b/keras/src/ops/core.py
@@ -254,7 +254,7 @@ def slice(inputs, start_indices, shape):
     inputs = np.zeros((5, 5))
     start_indices = np.array([3, 3])
     shape = np.array([2, 2])
-    inputs = keras.ops.slice(inputs, start_indices, updates)
+    inputs = keras.ops.slice(inputs, start_indices, shape)
     ```
 
     Args:
diff --git a/keras/src/trainers/trainer.py b/keras/src/trainers/trainer.py
index af5e3a3f7..0c7880db7 100644
--- a/keras/src/trainers/trainer.py
+++ b/keras/src/trainers/trainer.py
@@ -1,3 +1,4 @@
+import inspect
 import platform
 import warnings
 
@@ -25,6 +26,9 @@ class Trainer:
         self.steps_per_execution = 1
         # Can be set by callbacks in on_train_begin
         self._initial_epoch = None
+        self._compute_loss_has_training_arg = (
+            "training" in inspect.signature(self.compute_loss).parameters
+        )
 
     @traceback_utils.filter_traceback
     @tracking.no_automatic_dependency_tracking
@@ -262,6 +266,7 @@ class Trainer:
         y=None,
         y_pred=None,
         sample_weight=None,
+        training=True,
     ):
         """Compute the total loss, validate it, and return it.
 
@@ -276,7 +281,7 @@ class Trainer:
                 super().__init__(*args, **kwargs)
                 self.loss_tracker = metrics.Mean(name='loss')
 
-            def compute_loss(self, x, y, y_pred, sample_weight):
+            def compute_loss(self, x, y, y_pred, sample_weight, training=True):
                 loss = ops.means((y_pred - y) ** 2)
                 loss += ops.sum(self.losses)
                 self.loss_tracker.update_state(loss)
@@ -306,12 +311,15 @@ class Trainer:
             y: Target data.
             y_pred: Predictions returned by the model (output of `model(x)`)
             sample_weight: Sample weights for weighting the loss function.
+            training: Whether we are training or evaluating the model.
 
         Returns:
             The total loss as a scalar tensor, or `None` if no loss results
             (which is the case when called by `Model.test_step`).
         """
-        del x  # The default implementation does not use `x`.
+        # The default implementation does not use `x` or `training`.
+        del x
+        del training
         losses = []
         if self._compile_loss is not None:
             loss = self._compile_loss(y, y_pred, sample_weight)
@@ -331,6 +339,27 @@ class Trainer:
             total_loss = ops.sum(losses)
         return total_loss
 
+    def _compute_loss(
+        self,
+        x=None,
+        y=None,
+        y_pred=None,
+        sample_weight=None,
+        training=True,
+    ):
+        """Backwards compatibility wrapper for `compute_loss`.
+
+        This should be used instead `compute_loss` within `train_step` and
+        `test_step` to support overrides of `compute_loss` that may not have
+        the `training` argument, as this argument was added in Keras 3.3.
+        """
+        if self._compute_loss_has_training_arg:
+            return self.compute_loss(
+                x, y, y_pred, sample_weight, training=training
+            )
+        else:
+            return self.compute_loss(x, y, y_pred, sample_weight)
+
     def stateless_compute_loss(
         self,
         trainable_variables,
@@ -340,6 +369,7 @@ class Trainer:
         y=None,
         y_pred=None,
         sample_weight=None,
+        training=True,
     ):
         var_mapping = list(zip(self.trainable_variables, trainable_variables))
         var_mapping.extend(
@@ -349,11 +379,12 @@ class Trainer:
         with backend.StatelessScope(state_mapping=var_mapping) as scope:
             # Note that this is needed for the regularization loss, which need
             # the latest value of train/non-trainable variables.
-            loss = self.compute_loss(
+            loss = self._compute_loss(
                 x,
                 y,
                 y_pred,
                 sample_weight=sample_weight,
+                training=training,
             )
 
         # Update non trainable vars (may have been updated in compute_loss)


Test patch:
diff --git a/keras/src/initializers/random_initializers_test.py b/keras/src/initializers/random_initializers_test.py
index b58cf64ba..c0fec3812 100644
--- a/keras/src/initializers/random_initializers_test.py
+++ b/keras/src/initializers/random_initializers_test.py
@@ -147,6 +147,10 @@ class InitializersTest(testing.TestCase):
 
         self.run_class_serialization_test(initializer)
 
+        # Test legacy class_name
+        initializer = initializers.get("Orthogonal")
+        self.assertIsInstance(initializer, initializers.OrthogonalInitializer)
+
     def test_get_method(self):
         obj = initializers.get("glorot_normal")
         self.assertTrue(obj, initializers.GlorotNormal)
diff --git a/keras/src/layers/rnn/gru_test.py b/keras/src/layers/rnn/gru_test.py
index 529220803..d2c4f9656 100644
--- a/keras/src/layers/rnn/gru_test.py
+++ b/keras/src/layers/rnn/gru_test.py
@@ -286,3 +286,26 @@ class GRUTest(testing.TestCase, parameterized.TestCase):
             np.array([[0.11669192, 0.11669192], [0.28380975, 0.28380975]]),
             output,
         )
+
+    def test_legacy_implementation_argument(self):
+        sequence = np.arange(72).reshape((3, 6, 4)).astype("float32")
+        layer = layers.GRU(
+            3,
+            kernel_initializer=initializers.Constant(0.01),
+            recurrent_initializer=initializers.Constant(0.02),
+            bias_initializer=initializers.Constant(0.03),
+        )
+        config = layer.get_config()
+        config["implementation"] = 0  # Add legacy argument
+        layer = layers.GRU.from_config(config)
+        output = layer(sequence)
+        self.assertAllClose(
+            np.array(
+                [
+                    [0.5217289, 0.5217289, 0.5217289],
+                    [0.6371659, 0.6371659, 0.6371659],
+                    [0.39384964, 0.39384964, 0.3938496],
+                ]
+            ),
+            output,
+        )
diff --git a/keras/src/trainers/trainer_test.py b/keras/src/trainers/trainer_test.py
index a7ec7b9b5..a4b21e5f5 100644
--- a/keras/src/trainers/trainer_test.py
+++ b/keras/src/trainers/trainer_test.py
@@ -1387,6 +1387,7 @@ class TestTrainer(testing.TestCase, parameterized.TestCase):
 
     @pytest.mark.requires_trainable_backend
     def test_metric_update_in_compute_loss(self):
+        test_self = self
 
         class MyModel(keras.Model):
             def __init__(self):
@@ -1398,9 +1399,17 @@ class TestTrainer(testing.TestCase, parameterized.TestCase):
                 return self.dense(x)
 
             def compute_loss(
-                self, x=None, y=None, y_pred=None, sample_weight=None
+                self,
+                x=None,
+                y=None,
+                y_pred=None,
+                sample_weight=None,
+                training=True,
             ):
-                loss = super().compute_loss(x, y, y_pred, sample_weight)
+                test_self.assertTrue(training)
+                loss = super().compute_loss(
+                    x, y, y_pred, sample_weight, training
+                )
                 self.custom_metric.update_state(loss * 4)
                 return loss
 
@@ -1415,6 +1424,7 @@ class TestTrainer(testing.TestCase, parameterized.TestCase):
 
     @pytest.mark.requires_trainable_backend
     def test_fwd_pass_loss_presence_in_compute_loss(self):
+        test_self = self
 
         class MyModel(keras.Model):
             def __init__(self):
@@ -1426,9 +1436,17 @@ class TestTrainer(testing.TestCase, parameterized.TestCase):
                 return self.dense(x)
 
             def compute_loss(
-                self, x=None, y=None, y_pred=None, sample_weight=None
+                self,
+                x=None,
+                y=None,
+                y_pred=None,
+                sample_weight=None,
+                training=True,
             ):
-                loss = super().compute_loss(x, y, y_pred, sample_weight)
+                test_self.assertTrue(training)
+                loss = super().compute_loss(
+                    x, y, y_pred, sample_weight, training
+                )
                 self.custom_metric.update_state(sum(self.losses))
                 return loss
 
@@ -1439,6 +1457,75 @@ class TestTrainer(testing.TestCase, parameterized.TestCase):
         history = model.fit(x, y)
         self.assertGreater(history.history["custom"][0], 0.0)
 
+    @pytest.mark.requires_trainable_backend
+    def test_evaluate_with_custom_compute_loss(self):
+        test_self = self
+
+        class MyModel(keras.Model):
+            def __init__(self):
+                super().__init__()
+                self.custom_metric = keras.metrics.Mean(name="custom")
+                self.dense = keras.layers.Dense(2, activity_regularizer="l2")
+
+            def call(self, x):
+                return self.dense(x)
+
+            def compute_loss(
+                self,
+                x=None,
+                y=None,
+                y_pred=None,
+                sample_weight=None,
+                training=True,
+            ):
+                test_self.assertFalse(training)
+                loss = super().compute_loss(
+                    x, y, y_pred, sample_weight, training
+                )
+                self.custom_metric.update_state(loss * 4)
+                return loss
+
+        model = MyModel()
+        model.compile(optimizer="sgd", loss="mse")
+        x = np.ones((32, 4))
+        y = np.ones((32, 2)) * 2
+        logs = model.evaluate(x, y, return_dict=True)
+        self.assertAlmostEqual(logs["custom"], logs["loss"] * 4)
+
+    @pytest.mark.requires_trainable_backend
+    def test_compute_loss_no_training_backwards_compatibility(self):
+
+        class MyModel(keras.Model):
+            def __init__(self):
+                super().__init__()
+                self.custom_metric = keras.metrics.Mean(name="custom")
+                self.dense = keras.layers.Dense(2, activity_regularizer="l2")
+
+            def call(self, x):
+                return self.dense(x)
+
+            def compute_loss(
+                self,
+                x=None,
+                y=None,
+                y_pred=None,
+                sample_weight=None,
+            ):
+                loss = super().compute_loss(x, y, y_pred, sample_weight)
+                self.custom_metric.update_state(loss * 4)
+                return loss
+
+        model = MyModel()
+        model.compile(optimizer="sgd", loss="mse")
+        x = np.ones((32, 4))
+        y = np.ones((32, 2)) * 2
+        logs = model.evaluate(x, y, return_dict=True)
+        self.assertAlmostEqual(logs["custom"], logs["loss"] * 4)
+        history = model.fit(x, y)
+        self.assertAlmostEqual(
+            history.history["custom"][0], history.history["loss"][0] * 4
+        )
+
     @pytest.mark.requires_trainable_backend
     def test_loss_weights(self):
         epochs = 3