diff options
Diffstat (limited to 'tensorflow/contrib/mpi_collectives/__init__.py')
| -rw-r--r-- | tensorflow/contrib/mpi_collectives/__init__.py | 18 |
1 files changed, 10 insertions, 8 deletions
diff --git a/tensorflow/contrib/mpi_collectives/__init__.py b/tensorflow/contrib/mpi_collectives/__init__.py index 9ed16a6f07..52029cbc36 100644 --- a/tensorflow/contrib/mpi_collectives/__init__.py +++ b/tensorflow/contrib/mpi_collectives/__init__.py @@ -37,7 +37,7 @@ for detecting the running MPI configuration. Example: ```python -from tensorflow.contrib import mpi +import tensorflow.contrib.mpi_collectives as mpi # Use `mpi.Session` instead of `tf.Session` with mpi.Session() as session: @@ -48,8 +48,10 @@ with mpi.Session() as session: print("MPI Size:", session.run(mpi.size())) ``` -@@rank +@@init @@size +@@rank +@@local_rank ### Ring Allreduce and Allgather @@ -123,12 +125,12 @@ from __future__ import print_function import tensorflow as tf -from tensorflow.contrib.mpi_collectives.mpi_ops import size -from tensorflow.contrib.mpi_collectives.mpi_ops import rank -from tensorflow.contrib.mpi_collectives.mpi_ops import local_rank -from tensorflow.contrib.mpi_collectives.mpi_ops import allgather -from tensorflow.contrib.mpi_collectives.mpi_ops import _allreduce -from tensorflow.contrib.mpi_collectives.mpi_ops import init +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import init +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import size +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import rank +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import local_rank +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import allgather +from tensorflow.contrib.mpi_collectives.python.ops.mpi_ops import _allreduce def allreduce(tensor, average=True): |