Repository: keras-team/keras (python)

Issue:
keras adamw optimizer failed with callable parameters in TensorFlow2.16
When we were working on upgrading keras 2 to keras 3 in TensorFlow plugin, one of our adamw related unit test failed, which is a sub unit test using callable  lambda as learning_rate argument. We also found this ut failed in TensorFlow2.16 official docker image. The error log is :
![image](https://github.com/keras-team/keras/assets/25453568/3f212c38-f1d1-4dc0-95bd-f8f3d8c37916)


```python
"""Tests for adam optimizer with weight decay."""

import numpy as np
import keras
import tensorflow as tf
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import constant_op
from tensorflow.python.platform import test
from tensorflow.python.framework import test_util

from keras.src.optimizers import adamw

DATA_TYPES = [
    dtypes.float32
]

WEIGHT_DECAY = 0.1

def adamw_update_numpy(
    param, grad_t, slot_vars, learning_rate, beta_1, beta_2, epsilon, weight_decay, amsgrad
):
    """Numpy update function for AdamW."""
    lr, beta1, beta2, eps, wd = (
        v() if callable(v) else v
        for v in (learning_rate, beta_1, beta_2, epsilon, weight_decay)
    )
    t = slot_vars.get("t", 0) + 1
    lr_t = lr * np.sqrt(1 - beta2 ** t) / (1 - beta1 ** t)
    slot_vars["m"] = beta1 * slot_vars.get("m", 0) + (1 - beta1) * grad_t
    slot_vars["v"] = beta2 * slot_vars.get("v", 0) + (1 - beta2) * grad_t ** 2
    if amsgrad:
        slot_vars["v_hat"] = slot_vars.get("v_hat", 0)
        slot_vars["v_hat"] = np.maximum(slot_vars["v_hat"], slot_vars["v"])
        param_t = param * (1 - wd * lr) - lr_t * slot_vars["m"] / (np.sqrt(slot_vars["v_hat"]) + eps)
    else:
        param_t = param * (1 - wd * lr) - lr_t * slot_vars["m"] / (np.sqrt(slot_vars["v"]) + eps)
    slot_vars["t"] = t
    return param_t, slot_vars

class AdamWeightDecayOptimizerTest(test_util.TensorFlowTestCase):

    def doTestBasic(self, use_callable_params=False, do_sparse=False, do_amsgrad=False):
        for dtype in DATA_TYPES:
            # Initialize variables for numpy implementation.
            np_slot_vars0, np_slot_vars1 = {}, {}
            var0_np = np.array([1.0, 2.0], dtype=dtype.as_numpy_dtype)
            grads0_np = np.array([0.1, 0.1], dtype=dtype.as_numpy_dtype)
            var1_np = np.array([3.0, 4.0], dtype=dtype.as_numpy_dtype)
            grads1_np = np.array([0.01, 0.01], dtype=dtype.as_numpy_dtype)

            # Create Tensorflow variables.
            itex_var0 = tf.Variable(var0_np)
            itex_var1 = tf.Variable(var1_np)

            # Adapt callable parameters
            learning_rate = lambda: 0.01
            beta_1=lambda: 0.9
            beta_2=lambda: 0.999
            if not use_callable_params:
                learning_rate = learning_rate()
                beta_1 = beta_1()
                beta_2 = beta_2()

            # Adapt sparse
            if do_sparse:
                grads0_np_indices = np.array([0, 1], dtype=np.int32)
                grads0 = tf.IndexedSlices(
                    tf.constant(grads0_np), tf.constant(grads0_np_indices), tf.constant([2])
                )
                grads1_np_indices = np.array([0, 1], dtype=np.int32)
                grads1 = tf.IndexedSlices(
                    tf.constant(grads1_np), tf.constant(grads1_np_indices), tf.constant([2])
                )
            else:
                grads0 = constant_op.constant(grads0_np)
                grads1 = constant_op.constant(grads1_np)

            adamw_opt = adamw.AdamW(weight_decay=WEIGHT_DECAY, learning_rate=learning_rate, amsgrad=do_amsgrad)

            # Run 3 steps of the optimizer
            for _ in range(3):
                adamw_opt.apply_gradients(
                    zip([grads0, grads1], [itex_var0, itex_var1])
                )
                var0_np, np_slot_vars0 = adamw_update_numpy(
                    var0_np, grads0_np, np_slot_vars0, weight_decay=WEIGHT_DECAY, learning_rate=learning_rate, 
                    beta_1=beta_1, beta_2=beta_2, epsilon=1e-7, amsgrad=do_amsgrad)
                var1_np, np_slot_vars1 = adamw_update_numpy(
                    var1_np, grads1_np, np_slot_vars1, weight_decay=WEIGHT_DECAY, learning_rate=learning_rate, 
                    beta_1=beta_1, beta_2=beta_2, epsilon=1e-7, amsgrad=do_amsgrad)
                # Validate updated parameters
                self.assertAllCloseAccordingToType(itex_var0.numpy(), var0_np)
                self.assertAllCloseAccordingToType(itex_var1.numpy(), var1_np)


    def testCallableParamsAdamW(self):
        '''ResourceApplyAdamWithWeightDecay is a DPCPP op, don't have cpu registration 
            TODO: waiting for CPU registration of ResourceApplyAdamWithWeightDecay then enable
            this test case on CPU'''
        if not test.is_gpu_available():
            self.skipTest("No GPU available")
        self.doTestBasic(use_callable_params=True)


if __name__ == "__main__":
    test.main()
```


Pull request description:
Fix issue with using a callable as a learning rate
Closes #19411.

Patch:
diff --git a/keras/backend/tensorflow/distribution_lib.py b/keras/backend/tensorflow/distribution_lib.py
index ca0806b8e..b5ce7c1ad 100644
--- a/keras/backend/tensorflow/distribution_lib.py
+++ b/keras/backend/tensorflow/distribution_lib.py
@@ -23,13 +23,22 @@ def list_devices(device_type=None):
     Return:
         List of devices that are available for distribute computation.
     """
-    device_type = device_type.upper() if device_type else "CPU"
+    device_type = device_type.upper() if device_type else None
 
     # DTensor doesn't support getting global devices, even when knowing the
     # Mesh. Use TF API instead to get global devices. Coordinator service is
     # enabled by default with DTensor, so that list_logical_devices() returns
     # a list of global devices. More context can be found in b/254911601.
     tf_devices = tf.config.list_logical_devices(device_type=device_type)
+    cpu_devices = []
+    other_devices = []
+    for device in tf_devices:
+        if device.device_type.lower() == "cpu":
+            cpu_devices.append(device)
+        else:
+            other_devices.append(device)
+    if device_type is None:
+        tf_devices = other_devices if len(other_devices) > 0 else cpu_devices
     return [
         f"{device.device_type.lower()}:{device.name.split(':')[-1]}"
         for device in tf_devices
diff --git a/keras/backend/tensorflow/optimizer.py b/keras/backend/tensorflow/optimizer.py
index d9307b9cc..9524149b7 100644
--- a/keras/backend/tensorflow/optimizer.py
+++ b/keras/backend/tensorflow/optimizer.py
@@ -128,12 +128,15 @@ class TFOptimizer(KerasAutoTrackable, base_optimizer.BaseOptimizer):
     ):
         grads_and_vars = self._all_reduce_sum_gradients(grads_and_vars)
 
-        def apply_grad_to_update_var(var, grad):
+        def apply_grad_to_update_var(var, grad, learning_rate):
             return self.update_step(grad, var, learning_rate)
 
         for grad, var in grads_and_vars:
             distribution.extended.update(
-                var, apply_grad_to_update_var, args=(grad,), group=False
+                var,
+                apply_grad_to_update_var,
+                args=(grad, learning_rate),
+                group=False,
             )
 
     def _all_reduce_sum_gradients(self, grads_and_vars):
diff --git a/keras/optimizers/base_optimizer.py b/keras/optimizers/base_optimizer.py
index e18f3c1de..8709a95b9 100644
--- a/keras/optimizers/base_optimizer.py
+++ b/keras/optimizers/base_optimizer.py
@@ -567,7 +567,7 @@ class BaseOptimizer:
         ):
             return self._learning_rate(self.iterations)
         elif callable(self._learning_rate):
-            return self._learning_rate(self.iterations)
+            return self._learning_rate()
         return self._learning_rate
 
     def _filter_empty_gradients(self, grads, vars):


Test patch:
diff --git a/keras/optimizers/optimizer_test.py b/keras/optimizers/optimizer_test.py
index f9ab11897..437350586 100644
--- a/keras/optimizers/optimizer_test.py
+++ b/keras/optimizers/optimizer_test.py
@@ -243,3 +243,12 @@ class OptimizerTest(testing.TestCase):
         checkpoint.restore(save_path)
         pred = model.predict(x)
         self.assertAllClose(pred, ref_pred, atol=1e-5)
+
+    def test_callable_learning_rate(self):
+        v = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
+        grads = backend.convert_to_tensor([[1.0, 1.0], [1.0, 1.0]])
+        optimizer = optimizers.AdamW(learning_rate=lambda: 0.0001)
+        self.assertAllClose(optimizer.iterations, 0)
+        optimizer.apply_gradients([(grads, v)])
+        self.assertAllClose(v, [[1.0, 2.0], [3.0, 4.0]], atol=1e-4)
+        self.assertAllClose(optimizer.iterations, 1)