Repository: keras-team/keras (python)

Issue:
Assign doesn't return a value
When assigning a placeholder value to a KerasVariable I get returned a None type as there is no return in assignment, where as when assigning to a tf.Variable I get returned the proper assign operation.

This is the code in keras.Variable:
```
def _direct_assign(self, value):
        self._value.assign(tf.cast(value, self._value.dtype))
```

And here in the KerasVariable:
```
    def assign(self, value):
        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 "
                "the shape of the target value in "
                "`variable.assign(value)` must match. "
                f"variable.shape={self.value.shape}, "
                f"Received: value.shape={value.shape}. "
                f"Target variable: {self}"
            )
        if in_stateless_scope():
            scope = get_stateless_scope()
            scope.add_update((self, value))
        else:
            self._direct_assign(value)
```

Link to Keras implementation: https://github.com/keras-team/keras/blob/f9501f53fa15fe6286ca834618a6bcf10f86f7e8/keras/src/backend/common/variables.py#L224


Pull request description:
Return value on KerasVariable assignment
To keep compatibility with previous Tensorflow implementations, we can return variable value after assignment.

Specifically, this impacts porting Tensorflow Federated to Keras 3, see https://github.com/google-parfait/tensorflow-federated/blob/0ca2550e62ca7728e5e9a45edca1feedc6bc065f/tensorflow_federated/python/learning/models/functional.py#L528 which depends on returning the value after an `assign`.

Overhead should be minimal and I could see this being useful for finding the value after an `assign_add` or `assign_sub`.

Fix #19997.

Patch:
diff --git a/keras/src/backend/common/variables.py b/keras/src/backend/common/variables.py
index 34b69958f..e76930f0a 100644
--- a/keras/src/backend/common/variables.py
+++ b/keras/src/backend/common/variables.py
@@ -237,12 +237,13 @@ class KerasVariable:
             scope.add_update((self, value))
         else:
             self._direct_assign(value)
+        return value
 
     def assign_add(self, value):
-        self.assign(self + value)
+        return self.assign(self + value)
 
     def assign_sub(self, value):
-        self.assign(self - value)
+        return self.assign(self - value)
 
     @property
     def dtype(self):


Test patch:
diff --git a/keras/src/backend/common/variables_test.py b/keras/src/backend/common/variables_test.py
index ab40c905c..2b36fbda5 100644
--- a/keras/src/backend/common/variables_test.py
+++ b/keras/src/backend/common/variables_test.py
@@ -297,18 +297,36 @@ class VariableNumpyValueAndAssignmentTest(test_case.TestCase):
         v.assign(np.array([4, 5, 6]))
         self.assertAllClose(v.value, np.array([4, 5, 6]))
 
+    def test_variable_assign_return(self):
+        """Test assigning a new value and returning."""
+        v = backend.Variable(initializer=np.array([1, 2, 3]))
+        r = v.assign(np.array([4, 5, 6]))
+        self.assertAllClose(r, np.array([4, 5, 6]))
+
     def test_variable_assign_add(self):
         """Test the assign_add method on a variable."""
         v = backend.Variable(initializer=np.array([1, 2, 3]))
         v.assign_add(np.array([1, 1, 1]))
         self.assertAllClose(v.value, np.array([2, 3, 4]))
 
+    def test_variable_assign_add_return(self):
+        """Test assign_add a new value and returning."""
+        v = backend.Variable(initializer=np.array([1, 2, 3]))
+        r = v.assign_add(np.array([1, 1, 1]))
+        self.assertAllClose(r, np.array([2, 3, 4]))
+
     def test_variable_assign_sub(self):
         """Test the assign_sub method on a variable."""
         v = backend.Variable(initializer=np.array([2, 3, 4]))
         v.assign_sub(np.array([1, 1, 1]))
         self.assertAllClose(v.value, np.array([1, 2, 3]))
 
+    def test_variable_assign_sub_return(self):
+        """Test assign_sub a new value and returning."""
+        v = backend.Variable(initializer=np.array([2, 3, 4]))
+        r = v.assign_sub(np.array([1, 1, 1]))
+        self.assertAllClose(r, np.array([1, 2, 3]))
+
     def test_deferred_initialize_within_stateless_scope(self):
         """Test deferred init within a stateless scope."""
         with backend.StatelessScope():