INNER CODE UNIT · Python
_push_pull
bytedance/byteps · byteps/_keras/__init__.py:47
def _push_pull(self, gradients):
self._aggregated_gradients = True
if bps.size() > 1:
averaged_gradients = []
with tf.name_scope(self._name + "_Push_Pull") as scope:
for grad in gradients:
if grad is not None:
if self._sparse_as_dense and \
isinstance(grad, tf.IndexedSlices):
grad = tf.convert_to_tensor(grad)
avg_grad = bps.push_pull(grad, scope,
device_dense=self._device_dense,
device_sparse=self._device_sparse,
compression=self._compression)
averaged_gradients.append(avg_grad)
else:
averaged_gradients.append(None)
return averaged_gradients