Repository: keras-team/keras

Issue:
apply_gradients AttributeError: 'ResourceVariable' object has no attribute 'overwrite_with_gradient'
When I have a mix of tf.Variable and KerasVariables I get the following error:
```
--> 632     if v.overwrite_with_gradient:
    633         if self.gradient_accumulation_steps:
    634             # Utilize a stateless manner for JAX compatibility
    635             steps = self.gradient_accumulation_steps

AttributeError: 'ResourceVariable' object has no attribute 'overwrite_with_gradient'
```

I suspect this is because my list of variables is [KerasVariables] + [tf.Variables]
and the following line only checks the first in the list as to whether overwrite_with_gradient can be used?
https://github.com/keras-team/keras/blob/660da946d33ccbb53575f15a17abdfd725cd2086/keras/src/optimizers/base_optimizer.py#L675



Discussion:
Could you please provide some sample reproducible script to replicate the reported behavior. Thanks!
> Could you please provide some sample reproducible script to replicate the reported behavior. Thanks!

@sachinprasadhs Sure, please try this cut down simple example showing the problem:
```
import tensorflow as tf
from tensorflow import keras 

class MyModel(tf.keras.Model):
    def __init__(self):
        super().__init__()
        
        # Keras model Layers
        self.hidden_layers = [tf.keras.layers.Dense(32, activation='tanh') for _ in range(2)]
        self.output_layer = tf.keras.layers.Dense(1)
        
        # Custom variable
        self.my_var = tf.Variable(0.1, trainable=True, dtype=tf.float32, name="my_var")

    def call(self, inputs):
        x = inputs
        for layer in self.hidden_layers:
            x = layer(x)
        return self.output_layer(x)
    
data = np.array([
    [0.0,    10.4],
    [900.0,  21.1],
    [3900.0, 64.2],
]
)

model   = MyModel()
inputs  = data[:, 0:1]
outputs = data[:, 1:]

epochs = 1000
learning_rate = 0.005
optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

for epoch in range(epochs):
    with tf.GradientTape() as tp:
        y_pred = model(inputs)
        loss   = tf.reduce_mean(tf.square((outputs - y_pred)))
    
    params = model.trainable_variables + [model.my_var]
        
    gradients = tp.gradient(loss, params)
    optimizer.apply_gradients(zip(gradients, params))
    del tp
```


@andrewl36 
Try this one:

```python
import numpy as np
import tensorflow as tf

import keras


class MyModel(keras.Model):
    def __init__(self):
        super().__init__()

        # Keras model Layers
        self.hidden_layers = [
            keras.layers.Dense(32, activation="tanh") for _ in range(2)
        ]
        self.output_layer = keras.layers.Dense(1)

        # Custom variable
        self.my_var = self.add_weight(shape=(), dtype="float32", name="my_var")
        self.my_var.assign(0.1)

    def call(self, inputs):
        x = inputs
        for layer in self.hidden_layers:
            x = layer(x)
        return self.output_layer(x)


data = np.array(
    [
        [0.0, 10.4],
        [900.0, 21.1],
        [3900.0, 64.2],
    ]
)

model = MyModel()
inputs = data[:, 0:1]
outputs = data[:, 1:]
epochs = 1000
learning_rate = 0.005
optimizer = keras.optimizers.Adam(learning_rate=learning_rate)

pbar = keras.utils.Progbar(epochs)
for epoch in range(epochs):
    with tf.GradientTape() as tp:
        y_pred = model(inputs)
        loss = tf.reduce_mean(tf.square((outputs - y_pred)))
    gradients = tp.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    values = {"loss": loss.numpy()}
    pbar.add(1, values.items())

```

The key is to use `self.add_weight` for creating custom variables.

EDITED:
I have submitted a PR to fix this.



Patch:
diff --git a/keras/src/optimizers/base_optimizer.py b/keras/src/optimizers/base_optimizer.py
index b966c15bc17e..8717192fa84f 100644
--- a/keras/src/optimizers/base_optimizer.py
+++ b/keras/src/optimizers/base_optimizer.py
@@ -672,7 +672,7 @@ def _overwrite_variables_directly_with_gradients(self, grads, vars):
         """
         # Shortcut for `tf.Variable` because it doesn't have a
         # `overwrite_with_gradient` attr
-        if not hasattr(vars[0], "overwrite_with_gradient"):
+        if any(not hasattr(v, "overwrite_with_gradient") for v in vars):
             return grads, vars
 
         # Shallow copies
@@ -723,7 +723,11 @@ def _filter_empty_gradients(self, grads, vars):
             if filtered_grads[i] is None:
                 filtered_grads.pop(i)
                 v = filtered_vars.pop(i)
-                missing_grad_vars.append(v.path)
+                try:
+                    missing_grad_vars.append(v.path)
+                except AttributeError:
+                    # `tf.Variable` doesn't have `path` attr.
+                    missing_grad_vars.append(v.name)
 
         if not filtered_grads:
             raise ValueError("No gradients provided for any variable.")


Test patch:
diff --git a/keras/src/optimizers/optimizer_test.py b/keras/src/optimizers/optimizer_test.py
index 1ac35e63cdd6..7d661df9a3c0 100644
--- a/keras/src/optimizers/optimizer_test.py
+++ b/keras/src/optimizers/optimizer_test.py
@@ -401,3 +401,25 @@ def test_pickleable_optimizers(self, optimizer):
         reloaded = pickle.loads(pickle.dumps(optimizer))
 
         self.assertEqual(optimizer.get_config(), reloaded.get_config())
+
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="The tf.Variable test can only run with TensorFlow backend.",
+    )
+    def test_mixed_with_tf_variables(self):
+        import tensorflow as tf
+
+        v = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
+        grads = backend.convert_to_tensor([[1.0, 1.0], [1.0, 1.0]])
+        tf_v = tf.Variable([[1.0, 2.0], [3.0, 4.0]])
+        tf_grads = backend.convert_to_tensor([[1.0, 1.0], [1.0, 1.0]])
+        optimizer = optimizers.Adam(learning_rate=1.0)
+        optimizer.apply_gradients([(grads, v), (tf_grads, tf_v)])
+        self.assertAllClose(optimizer.iterations, 1)
+
+        # Test with no grads
+        with self.assertWarnsRegex(
+            UserWarning, "Gradients do not exist for variables"
+        ):
+            optimizer.apply_gradients([(grads, v), (None, tf_v)])
+            self.assertAllClose(optimizer.iterations, 2)