Fix mnist sgd jenkins tests on master (#4168)

python/ray/experimental/sgd/mnist_example.py

@@ -24,7 +24,7 @@ from ray.tune import run_experiments
 from ray.tune.examples.tune_mnist_ray import deepnn
 from ray.experimental.sgd.model import Model
 from ray.experimental.sgd.sgd import DistributedSGD
-import ray.experimental.tf_utils
+import ray.experimental.tf_utils as ray_tf_utils
 
 parser = argparse.ArgumentParser()
 parser.add_argument("--redis-address", default=None, type=str)
@@ -67,7 +67,7 @@ class MNISTModel(Model):
             tf.nn.softmax_cross_entropy_with_logits(
                 labels=self.y_, logits=y_conv))
         self.optimizer = tf.train.AdamOptimizer(1e-4)
-        self.variables = ray.experimental.tfutils.TensorFlowVariables(
+        self.variables = ray_tf_utils.TensorFlowVariables(
             self.loss, tf.get_default_session())
 
         # For evaluating test accuracy

python/ray/experimental/sgd/modified_allreduce.py

@@ -27,7 +27,6 @@ import re
 from six.moves import xrange  # pylint: disable=redefined-builtin
 import tensorflow as tf
 
-from tensorflow.contrib import nccl
 from tensorflow.contrib.all_reduce.python import all_reduce
 
 logger = logging.getLogger(__name__)
@@ -309,7 +308,8 @@ def sum_grad_and_var_all_reduce(grad_and_vars,
         #   ((grad0_gpu0, var0_gpu0), ... , (grad0_gpuN, var0_gpuN))
         scaled_grads = [g for g, _ in grad_and_vars]
         if alg == 'nccl':
-            summed_grads = nccl.all_sum(scaled_grads)
+            from tensorflow.python.ops import nccl_ops
+            summed_grads = nccl_ops.all_sum(scaled_grads)
         elif alg == 'simple':
             summed_grads = build_reduce_sum(scaled_grads)
         elif alg == 'trivial':

python/ray/experimental/sgd/tfbench/test_model.py

@@ -6,7 +6,7 @@ import tensorflow as tf
 
 from tfbench import model_config
 from ray.experimental.sgd.model import Model
-import ray.experimental.tf_utils
+import ray.experimental.tf_utils as ray_tf_utils
 
 
 class MockDataset():
@@ -47,7 +47,7 @@ class TFBenchModel(Model):
         self.loss = tf.reduce_mean(loss, name='xentropy-loss')
         self.optimizer = tf.train.GradientDescentOptimizer(1e-6)
 
-        self.variables = ray.experimental.tf_utils.TensorFlowVariables(
+        self.variables = ray_tf_utils.TensorFlowVariables(
             self.loss, tf.get_default_session())
 
     def get_loss(self):