Fix assigning a value to a variable within an autocast scope. (#21864)

Previously `assign` would incorrectly cast the value to assign to the autocast dtype instead of the true dtype of the variable.

Because on JAX and OpenVino variables are just a reference to an array, this would cause the variable value to change dtypes.

Author: hertschuh

keras/src/backend/common/variables.py

@@ -276,7 +276,7 @@ def value(self):
         return self._maybe_autocast(self._value)
 
     def assign(self, value):
-        value = self._convert_to_tensor(value, dtype=self.dtype)
+        value = self._convert_to_tensor(value, dtype=self._dtype)
         if not shape_equal(value.shape, self.shape):
             raise ValueError(
                 "The shape of the target variable and "

keras/src/backend/common/variables_test.py

@@ -176,8 +176,8 @@ def test_trainable_setter(self):
             v.trainable = False
             self.assertFalse(v._value.requires_grad)
 
-    def test_autocasting(self):
-        """Tests autocasting of float variables."""
+    def test_autocasting_float(self):
+        # Tests autocasting of float variables
         v = backend.Variable(
             initializer=initializers.RandomNormal(),
             shape=(2, 2),
@@ -191,6 +191,33 @@ def test_autocasting(self):
             )
         self.assertEqual(backend.standardize_dtype(v.value.dtype), "float32")
 
+    def test_autocasting_float_assign(self):
+        # Tests assigning value to variable within an autocast scope
+        v = backend.Variable(
+            initializer=initializers.RandomNormal(),
+            shape=(2, 2),
+            dtype="float32",
+        )
+        self.assertEqual(v.dtype, "float32")
+        self.assertEqual(backend.standardize_dtype(v.value.dtype), "float32")
+
+        # Assign float16 value within float16 scope
+        with AutocastScope("float16"):
+            self.assertEqual(
+                backend.standardize_dtype(v.value.dtype), "float16"
+            )
+            v.assign(ops.ones((2, 2), "float16"))
+        self.assertEqual(backend.standardize_dtype(v.value.dtype), "float32")
+
+        # Assign float32 value within float16 scope
+        with AutocastScope("float16"):
+            self.assertEqual(
+                backend.standardize_dtype(v.value.dtype), "float16"
+            )
+            v.assign(ops.zeros((2, 2), "float32"))
+        self.assertEqual(backend.standardize_dtype(v.value.dtype), "float32")
+
+    def test_autocasting_int(self):
         # Test non-float variables are not affected
         v = backend.Variable(
             initializer=initializers.Ones(),
@@ -204,6 +231,7 @@ def test_autocasting(self):
         with AutocastScope("float16"):
             self.assertEqual(backend.standardize_dtype(v.value.dtype), "int32")
 
+    def test_autocasting_float_with_autocast_off(self):
         # Test autocast argument
         v = backend.Variable(
             initializer=initializers.RandomNormal(),