keras/src/backend/tensorflow/optimizer.py/TFOptimizer/_backend_update_step
class TFOptimizer:    def _backend_update_step(self, grads, trainable_variables, learning_rate):
        trainable_variables = [
            v.value if isinstance(v, backend.Variable) else v
            for v in trainable_variables
        ]
        grads_and_vars = list(zip(grads, trainable_variables))
        grads_and_vars = self._all_reduce_sum_gradients(grads_and_vars)
        tf.__internal__.distribute.interim.maybe_merge_call(
            self._distributed_tf_update_step,
            self._distribution_strategy,
            grads_and_vars,
            learning_rate,
        )