From 07426bcdd275276403fd3dbd3e0773699afb7a76 Mon Sep 17 00:00:00 2001 From: Josh Romero Date: Thu, 24 Sep 2020 09:39:24 -0700 Subject: [PATCH] Expose name argument to TF allreduce operation. Signed-off-by: Josh Romero --- horovod/tensorflow/__init__.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/horovod/tensorflow/__init__.py b/horovod/tensorflow/__init__.py index 76c8d4d7ae..c3f711ef31 100644 --- a/horovod/tensorflow/__init__.py +++ b/horovod/tensorflow/__init__.py @@ -52,7 +52,8 @@ def allreduce(tensor, average=None, device_dense='', device_sparse='', compression=Compression.none, op=None, - prescale_factor=1.0, postscale_factor=1.0): + prescale_factor=1.0, postscale_factor=1.0, + name=None): """Perform an allreduce on a tf.Tensor or tf.IndexedSlices. This function performs a bandwidth-optimal ring allreduce on the input @@ -79,6 +80,7 @@ def allreduce(tensor, average=None, device_dense='', device_sparse='', Defaults to Average if None is given. prescale_factor: Multiplicative factor to scale tensor before allreduce. postscale_factor: Multiplicative factor to scale tensor after allreduce. + name: A name of the allreduce operation Returns: A tensor of the same shape and type as `tensor`, summed across all @@ -116,7 +118,8 @@ def allreduce(tensor, average=None, device_dense='', device_sparse='', tensor_compressed, ctx = compression.compress(tensor) summed_tensor_compressed = _allreduce(tensor_compressed, op=op, prescale_factor=prescale_factor, - postscale_factor=postscale_factor) + postscale_factor=postscale_factor, + name=name) summed_tensor = compression.decompress(summed_tensor_compressed, ctx) if op == Adasum: if 'CPU' not in tensor.device and gpu_available('tensorflow'):