Repository: keras-team/keras (python)

Issue:
TensorFlow GPU - Fix `keras/callbacks/backup_and_restore_callback_test.py`
Fix TF GPU Test error in keras/callbacks/backup_and_restore_callback_test.py and update TODO in https://github.com/keras-team/keras/blob/master/keras/kokoro/github/ubuntu/gpu/build.sh#L41

https://source.cloud.google.com/results/invocations/63fcc5e7-8577-43b9-912c-9f0ceb3413d2/targets/keras%2Fgithub%2Fubuntu%2Fgpu%2Ftensorflow%2Fpresubmit/log

```
______________ BackupAndRestoreCallbackTest.test_best_case_epoch _______________

self = <keras.callbacks.backup_and_restore_callback_test.BackupAndRestoreCallbackTest testMethod=test_best_case_epoch>

    @pytest.mark.requires_trainable_backend
    def test_best_case_epoch(self):
        temp_dir = self.get_temp_dir()
        backup_dir = file_utils.join(temp_dir, "subdir")
        self.assertFalse(file_utils.exists(backup_dir))

        model = self.make_model()
        self.assertEqual(int(model.layers[0].counter.value), 0)
        cbk = callbacks.BackupAndRestore(
            backup_dir=backup_dir, save_freq="epoch"
        )

        x_train = np.random.random((10, 3))
        y_train = np.random.random((10, 1))

        try:
>           model.fit(
                x_train,
                y_train,
                batch_size=4,
                callbacks=[
                    cbk,
                    InterruptingCallback(steps_int=None, epoch_int=2),
                ],
                epochs=6,
                verbose=0,
            )

keras/callbacks/backup_and_restore_callback_test.py:126:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
keras/utils/traceback_utils.py:114: in error_handler
    return fn(*args, **kwargs)
keras/backend/tensorflow/trainer.py:322: in fit
    logs = self.train_function(iterator)
/tmpfs/venv/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py:153: in error_handler
    raise e.with_traceback(filtered_tb) from None
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

op_name = '__inference_one_step_on_iterator_8630', num_outputs = 2
inputs = [<tf.Tensor: shape=(), dtype=resource, value=<ResourceHandle(name="Resource-266-at-0x1800fd50", device="/job:localhost...evice:GPU:0", container="Anonymous", type="tensorflow::Var", dtype and shapes : "[ DType enum: 1, Shape: [] ]")>>, ...]
attrs = ('executor_type', '', 'config_proto', b'\n\x07\n\x03CPU\x10\x02\n\x07\n\x03GPU\x10\x042\x13*\x070,1,2,3J\x08\n\x00\n\x00\n\x00\n\x008\x01\x82\x01\x00')
ctx = <tensorflow.python.eager.context.Context object at 0x7f4fb807df10>
name = None

    def quick_execute(op_name, num_outputs, inputs, attrs, ctx, name=None):
      """Execute a TensorFlow operation.

      Args:
        op_name: Name of the TensorFlow operation (see REGISTER_OP in C++ code) to
          execute.
        num_outputs: The number of outputs of the operation to fetch. (Explicitly
          provided instead of being inferred for performance reasons).
        inputs: A list of inputs to the operation. Each entry should be a Tensor, or
          a value which can be passed to the Tensor constructor to create one.
        attrs: A tuple with alternating string attr names and attr values for this
          operation.
        ctx: The value of context.context().
        name: Customized name for the operation.

      Returns:
        List of output Tensor objects. The list is empty if there are no outputs

      Raises:
        An exception on error.
      """
      device_name = ctx.device_name
      # pylint: disable=protected-access
      try:
        ctx.ensure_initialized()
>       tensors = pywrap_tfe.TFE_Py_Execute(ctx._handle, device_name, op_name,
                                            inputs, attrs, num_outputs)
E                                           tensorflow.python.framework.errors_impl.InvalidArgumentError: Graph execution error:
E
E                                           Detected at node StatefulPartitionedCall defined at (most recent call last):

E                                             File "/tmpfs/src/github/keras/keras/callbacks/backup_and_restore_callback_test.py", line 126, in test_best_case_epoch
E
E                                             File "/tmpfs/src/github/keras/keras/utils/traceback_utils.py", line 114, in error_handler
E
E                                             File "/tmpfs/src/github/keras/keras/backend/tensorflow/trainer.py", line 322, in fit
E
E                                             File "/tmpfs/src/github/keras/keras/backend/tensorflow/trainer.py", line 117, in one_step_on_iterator
E
E                                           Trying to access resource canary_layer/variable_96/261 (defined @ /tmpfs/src/github/keras/keras/backend/tensorflow/core.py:30) located in device /job:localhost/replica:0/task:0/device:CPU:0 from device /job:localhost/replica:0/task:0/device:GPU:0
E                                            Cf. https://www.tensorflow.org/xla/known_issues#tfvariable_on_a_different_device
E                                           	 [[{{node StatefulPartitionedCall}}]] [Op:__inference_one_step_on_iterator_8630]

/tmpfs/venv/lib/python3.9/site-packages/tensorflow/python/eager/execute.py:53: InvalidArgumentError
----------------------------- Captured stderr call -----------------------------
2023-10-06 12:59:02.533913: W tensorflow/core/framework/op_kernel.cc:1839] OP_REQUIRES failed at xla_ops.cc:512 : INVALID_ARGUMENT: Trying to access resource canary_layer/variable_96/261 (defined @ /tmpfs/src/github/keras/keras/backend/tensorflow/core.py:30) located in device /job:localhost/replica:0/task:0/device:CPU:0 from device /job:localhost/replica:0/task:0/device:GPU:0
 Cf. https://www.tensorflow.org/xla/known_issues#tfvariable_on_a_different_device
_______________ BackupAndRestoreCallbackTest.test_best_case_step _______________
```

Pull request description:
Fix BackupRestore Test for TensorFlow GPU CI
Fixes #18565 

Patch:
diff --git a/keras/backend/tensorflow/numpy.py b/keras/backend/tensorflow/numpy.py
index af4c46138..c9742d983 100644
--- a/keras/backend/tensorflow/numpy.py
+++ b/keras/backend/tensorflow/numpy.py
@@ -810,7 +810,7 @@ def take(x, indices, axis=None):
             )
             return tfnp.take(x, tf.sparse.to_dense(indices), axis=axis)
         return tf.nn.safe_embedding_lookup_sparse(
-            embedding_weights=x,
+            embedding_weights=tf.convert_to_tensor(x),
             sparse_ids=tf.sparse.expand_dims(indices, axis=-1),
             default_id=0,
         )
diff --git a/keras/kokoro/github/ubuntu/gpu/build.sh b/keras/kokoro/github/ubuntu/gpu/build.sh
index 22f6a8313..7689acd3e 100644
--- a/keras/kokoro/github/ubuntu/gpu/build.sh
+++ b/keras/kokoro/github/ubuntu/gpu/build.sh
@@ -38,11 +38,9 @@ then
 
    # TODO: keras/layers/merging/merging_test.py::MergingLayersTest::test_sparse_dot_2d Fatal Python error: Aborted
    # TODO: Embedding test failure
-   # TODO: Backup and Restore fails
    pytest keras --ignore keras/applications \
                --ignore keras/layers/merging/merging_test.py \
                --ignore keras/layers/core/embedding_test.py \
-               --ignore keras/callbacks/backup_and_restore_callback_test.py \
                --cov=keras
 fi
 
diff --git a/keras/kokoro/github/ubuntu/gpu/jax/continuous.cfg b/keras/kokoro/github/ubuntu/gpu/jax/continuous.cfg
new file mode 100644
index 000000000..a414d8e72
--- /dev/null
+++ b/keras/kokoro/github/ubuntu/gpu/jax/continuous.cfg
@@ -0,0 +1,8 @@
+build_file: "keras/keras/kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
diff --git a/keras/kokoro/github/ubuntu/gpu/jax/presubmit.cfg b/keras/kokoro/github/ubuntu/gpu/jax/presubmit.cfg
new file mode 100644
index 000000000..c87eb6547
--- /dev/null
+++ b/keras/kokoro/github/ubuntu/gpu/jax/presubmit.cfg
@@ -0,0 +1,13 @@
+build_file: "keras/keras/kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "jax"
+}
\ No newline at end of file
diff --git a/keras/kokoro/github/ubuntu/gpu/torch/continuous.cfg b/keras/kokoro/github/ubuntu/gpu/torch/continuous.cfg
new file mode 100644
index 000000000..a414d8e72
--- /dev/null
+++ b/keras/kokoro/github/ubuntu/gpu/torch/continuous.cfg
@@ -0,0 +1,8 @@
+build_file: "keras/keras/kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
diff --git a/keras/kokoro/github/ubuntu/gpu/torch/presubmit.cfg b/keras/kokoro/github/ubuntu/gpu/torch/presubmit.cfg
new file mode 100644
index 000000000..779550ec6
--- /dev/null
+++ b/keras/kokoro/github/ubuntu/gpu/torch/presubmit.cfg
@@ -0,0 +1,13 @@
+build_file: "keras/keras/kokoro/github/ubuntu/gpu/build.sh"
+
+action {
+  define_artifacts {
+    regex: "**/sponge_log.log"
+    regex: "**/sponge_log.xml"
+  }
+}
+
+env_vars: {
+   key: "KERAS_BACKEND"
+   value: "torch"
+}
\ No newline at end of file
diff --git a/keras/metrics/reduction_metrics.py b/keras/metrics/reduction_metrics.py
index 4199b9c08..8451b1959 100644
--- a/keras/metrics/reduction_metrics.py
+++ b/keras/metrics/reduction_metrics.py
@@ -197,7 +197,7 @@ class MeanMetricWrapper(Mean):
 
     def get_config(self):
         base_config = super().get_config()
-        config = {"fn": serialization_lib.serialize_keras_object(self.fn)}
+        config = {"fn": serialization_lib.serialize_keras_object(self._fn)}
         config.update(serialization_lib.serialize_keras_object(self._fn_kwargs))
         return {**base_config, **config}
 


Test patch:
diff --git a/keras/callbacks/backup_and_restore_callback_test.py b/keras/callbacks/backup_and_restore_callback_test.py
index 38eefe783..c6bd33c6f 100644
--- a/keras/callbacks/backup_and_restore_callback_test.py
+++ b/keras/callbacks/backup_and_restore_callback_test.py
@@ -32,7 +32,7 @@ class CanaryLayer(layers.Layer):
     def __init__(self):
         super().__init__()
         self.counter = self.add_weight(
-            shape=(), initializer="zeros", dtype="int32", trainable=False
+            shape=(), initializer="zeros", dtype="float32", trainable=False
         )
 
     def call(self, x):
diff --git a/keras/metrics/reduction_metrics_test.py b/keras/metrics/reduction_metrics_test.py
index 6fe94654d..b581c2620 100644
--- a/keras/metrics/reduction_metrics_test.py
+++ b/keras/metrics/reduction_metrics_test.py
@@ -2,6 +2,7 @@ import numpy as np
 
 from keras import testing
 from keras.metrics import reduction_metrics
+from keras.saving import register_keras_serializable
 
 
 class SumTest(testing.TestCase):
@@ -58,6 +59,8 @@ class MeanTest(testing.TestCase):
         self.assertAllClose(result, 3.0, atol=1e-3)
 
 
+# How users would register a custom function or class to use with MeanMetricWrapper.
+@register_keras_serializable(package="test", name="mse")
 def mse(y_true, y_pred):
     return (y_true - y_pred) ** 2
 
@@ -70,7 +73,14 @@ class MetricWrapperTest(testing.TestCase):
         self.assertEqual(mse_obj.name, "mse")
         self.assertEqual(len(mse_obj.variables), 2)
         self.assertEqual(mse_obj._dtype, "float32")
-        # TODO: Check save and restore config
+        # Check save and restore config
+        mse_obj2 = reduction_metrics.MeanMetricWrapper.from_config(
+            mse_obj.get_config()
+        )
+        self.assertEqual(mse_obj2.name, "mse")
+        self.assertEqual(len(mse_obj2.variables), 2)
+        self.assertEqual(mse_obj2._dtype, "float32")
+        self.assertTrue("fn" in mse_obj2.get_config())
 
     def test_unweighted(self):
         mse_obj = reduction_metrics.MeanMetricWrapper(