# Buggy code
def apply_gradients(self, grads_and_vars, name: Optional[str] = None, **kwargs):
    """Apply gradients to variables for each optimizer.

    On the first call to `apply_gradients()`, compute the mapping from variables to
    optimizers and cache it in the `self.var_opt_mapping` dict for serialization and
    faster access.
    """

    if self.var_opt_mapping is None:
        # Convert `grads_and_vars` to list so we can iterate multiple times over it
        grads_and_vars = list(grads_and_vars)
        self._compute_var_opt_mapping(grads_and_vars)

    # Split gradients and variables into a separate list for each optimizer
    grad_var_lists = [[] for _ in range(len(self.pred_opt_pairs) + 1)]
    for grad, var in grads_and_vars:
        if var.name in self.var_opt_mapping:
            grad_var_lists[self.var_opt_mapping[var.name]].append((grad, var))

    # Apply gradients to each optimizer
    with tf.name_scope(self._name):
        train_ops = [
            optimizer.apply_gradients(opt_grads_and_vars, **kwargs)
            for optimizer, opt_grads_and_vars in zip(self.optimizers, grad_var_lists)
        ]

        return tf.group(*train_ops, name=name or "train_with_group")


# Traceback (IndexError)
WARNING:tensorflow:There is non-GPU devices in `tf.distribute.Strategy`, not using nccl allreduce.
distributed training:  False
Train on 60000 samples
60000/60000 [==============================] - 4s 61us/sample - loss: 8.2390
Successfully fitted model

distributed training:  True
Train on 60000 samples
INFO:tensorflow:Error reported to Coordinator: list index out of range
Traceback (most recent call last):
File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/training/coordinator.py", line 297, in stop_on_exception
yield
File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 190, in _call_for_each_replica
**merge_kwargs)
File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/keras/optimizer_v2/optimizer_v2.py", line 446, in _distributed_apply
ds_reduce_util.ReduceOp.SUM, grads_and_vars)
File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/distribute_lib.py", line 1481, in batch_reduce_to
return self._batch_reduce_to(reduce_op, value_destination_pairs)
File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/mirrored_strategy.py", line 707, in _batch_reduce_to
value_destination_pairs)
File "/usr/local/lib/python3.7/dist-packages/tensorflow_core/python/distribute/cross_device_ops.py", line 317, in batch_reduce
value_destination_pairs[0][0].values) == 1:
IndexError: list index out of range
32/60000 [..............................] - ETA: 10:32Exception raised:
list index out of range

# Fixed code
def apply_gradients(self, grads_and_vars, name: Optional[str] = None, **kwargs):
    """Apply gradients to variables for each optimizer.

    On the first call to `apply_gradients()`, compute the mapping from variables to
    optimizers and cache it in the `self.var_opt_mapping` dict for serialization and
    faster access.
    """

    if self.var_opt_mapping is None:
        # Convert `grads_and_vars` to list so we can iterate multiple times over it
        grads_and_vars = list(grads_and_vars)
        self._compute_var_opt_mapping(grads_and_vars)

    # Split gradients and variables into a separate list for each optimizer
    grad_var_lists = [[] for _ in range(len(self.pred_opt_pairs) + 1)]
    for grad, var in grads_and_vars:
        if var.name in self.var_opt_mapping:
            grad_var_lists[self.var_opt_mapping[var.name]].append((grad, var))

    with tf.init_scope():
        for optimizer, opt_grads_and_vars in zip(self.optimizers, grad_var_lists):
            optimizer._create_slots([v for (_, v) in grads_and_vars])

    return tf.distribute.get_replica_context().merge_call(
        self._apply_gradients, args=(grad_var_lists, name), kwargs=kwargs
    )