keras/src/backend/tensorflow/optimizer.py/TFOptimizer/_distributed_tf_update_step
class TFOptimizer:    def _distributed_tf_update_step(
        self, distribution, grads_and_vars, learning_rate
    ):
        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, learning_rate),
                group=False,
            )