From fd48a4924cfcba4ebd5d9bb37060fb5c353a44c6 Mon Sep 17 00:00:00 2001
From: Qianli Scott Zhu <scottzhu@google.com>
Date: Tue, 24 Oct 2023 00:07:57 -0700
Subject: [PATCH] Fix numpy conversion for tf MirroredVariable. (#18675)

* Fix numpy conversion for tf MirroredVariable.

This is a corner case that np.array requires the return type of __array__
to be an array, and scalar will fail the type check.

* Fix format.

* Fix jax backend.

* Fix backend.
---
 integration_tests/tf_distribute_training_test.py | 12 ++++++++++--
 keras/backend/common/variables.py                |  6 +++++-
 keras/backend/common/variables_test.py           | 15 +++++++++++++++
 3 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/integration_tests/tf_distribute_training_test.py b/integration_tests/tf_distribute_training_test.py
index acff6a624f7..904b8e9238e 100644
--- a/integration_tests/tf_distribute_training_test.py
+++ b/integration_tests/tf_distribute_training_test.py
@@ -7,6 +7,7 @@
 from keras import metrics
 from keras import models
 from keras import optimizers
+from keras.callbacks import LearningRateScheduler
 
 
 def test_model_fit():
@@ -31,6 +32,8 @@ def test_model_fit():
         outputs = layers.Dense(16)(x)
         model = models.Model(inputs, outputs)
 
+    callbacks = [LearningRateScheduler(lambda _: 0.1)]
+
     model.summary()
 
     x = np.random.random((5000, 100))
@@ -49,7 +52,12 @@ def test_model_fit():
             jit_compile=False,
         )
         history = model.fit(
-            x, y, batch_size=batch_size, epochs=epochs, validation_split=0.2
+            x,
+            y,
+            batch_size=batch_size,
+            epochs=epochs,
+            validation_split=0.2,
+            callbacks=callbacks,
         )
 
     print("History:")
@@ -59,7 +67,7 @@ def test_model_fit():
     with strategy.scope():
         dataset = tf.data.Dataset.from_tensor_slices((x, y)).batch(batch_size)
         dataset = strategy.experimental_distribute_dataset(dataset)
-        history = model.fit(dataset, epochs=epochs)
+        history = model.fit(dataset, epochs=epochs, callbacks=callbacks)
 
     print("History:")
     print(history.history)
diff --git a/keras/backend/common/variables.py b/keras/backend/common/variables.py
index b5e171bb5bc..2706ad2f261 100644
--- a/keras/backend/common/variables.py
+++ b/keras/backend/common/variables.py
@@ -179,7 +179,11 @@ def __getitem__(self, idx):
         return self.value.__getitem__(idx)
 
     def __array__(self, dtype=None):
-        return self.value.__array__(dtype)
+        # We can't directly use self.value.__array__ here because of scalar.
+        # Numpy require this method to return as array like object. In the case
+        # of scalar, it will fail the type checking from numpy. We need to
+        # return a 0d array via numpy.
+        return np.asarray(self.value.__array__(dtype))
 
     def __bool__(self):
         raise TypeError("A Keras Variable cannot be used as a boolean.")
diff --git a/keras/backend/common/variables_test.py b/keras/backend/common/variables_test.py
index d52ccce1209..063eb9248b0 100644
--- a/keras/backend/common/variables_test.py
+++ b/keras/backend/common/variables_test.py
@@ -211,6 +211,21 @@ def test_variable_numpy(self):
         self.assertIsInstance(v.numpy(), np.ndarray)
         self.assertAllClose(v.numpy(), np.array([1, 2, 3]))
 
+    @pytest.mark.skipif(
+        backend.backend() != "tf",
+        reason="Tests for MirroredVariable under tf backend",
+    )
+    def test_variable_numpy_scalar(self):
+        from keras.utils.module_utils import tensorflow as tf
+
+        strategy = tf.distribute.MirroredStrategy(["cpu:0", "cpu:1"])
+        with strategy.scope():
+            v = backend.Variable(initializer=0.0)
+
+        np_value = backend.convert_to_numpy(v)
+        self.assertIsInstance(np_value, np.ndarray)
+        self.assertAllClose(np_value, 0.0)
+
     def test_variable_value(self):
         """Test retrieving the value of a variable."""
         v = backend.Variable(initializer=np.array([1, 2, 3]))