def _broadcast_nccl(self):
        """Sum gradients across devices using NCCL ops (fast path)."""
        from tensorflow.python.ops import nccl_ops # pylint: disable=no-name-in-module
        for all_vars in zip(*[device.grad_clean.keys() for device in self._devices.values()]):
            if any(x.shape.num_elements() > 0 for x in all_vars):
                all_grads = [device.grad_clean[var] for device, var in zip(self._devices.values(), all_vars)]
                all_grads = nccl_ops.all_sum(all_grads)
                for device, var, grad in zip(self._devices.values(), all_vars, all_grads):
                    device.grad_clean[var] = grad
Sum gradients across devices using NCCL ops (fast path).
converters/stylegan2ada_tf_official/dnnlib/tflib/optimizer.py