Repository: horovod/horovod
Pull request: Initial support for TensorFlow 2.0
* Support for TensorFlow 2.0 in `horovod.tensorflow`
* Replace RMSProp with Adam (RMSProp doesn't scale)
* Minor example bugfixes

Fixes #1177

Problem statement:
Horovod GradientTape performance for Tensorflow 
**Environment:**
1. Framework: Tensorflow
2. Framework version: 1.12
3. Horovod version:0.16.4
4. MPI version: Mvapich-GDR 2.3.1
5. CUDA version: 9.2
6. NCCL version:
7. Python version: 3.6
8. OS and version: Red Hat Enterprise Linux Server, 7.5 (Maipo)
9. GCC version: xl/2019.02.07
10. Architecture: Power9, V100 GPU


**Your question:**
I have modified tensorflow_mnist_eager.py script in the examples to print images per sec. 
Batch size: 32
#GPUs    Perf(images/sec)
1               2216
2              425
4              640 

It looks like there is an initial overhead of distributing the DNN training.
Is it the expected behavior?

Hints from issue discussion:
This is not expected.  @tgaddair, could you take a look?
@aj-prime, by the way, any reason you're not using NCCL?
@alsrgv No specific reason
@aj-prime, NCCL will give you better performance on GPU compared to MPI.  That said, I do see a slowdown of Eager TF compared to a regular TF even with NCCL in my environment.
@alsrgv I am trying to setup horovod with NCCL. Is the slowdown as severe as MV2-GDR?
@aj-prime, I have not tried MV2-GDR, so not sure.  What kind of performance are you seeing with Graph TensorFlow MNIST example?
@alsrgv For tensorflow_mnist.py script with batch size 32, I am getting following numbers 
1 GPU:   11292
2 GPUs: 12300
4 GPUs: 20080
OK, that's much better.  We'll look into the eager mode performance.
Hey @aj-prime, can you try running [tensorflow_synthetic_benchmark.py](https://github.com/horovod/horovod/blob/master/examples/tensorflow_synthetic_benchmark.py) with `--eager` and without?  That uses ResNet50, which might provide more interesting data.
hello @tgaddair 

I ran tensorflow_synthetic_benchmark.py in the graph and eager mode. Here are the results 

Graph Mode
1 GPU: 336.7 +-0.3
2 GPUS: 574.3 +-3.2
4 GPUs: 1090.5 +-13.1

Eager Mode
1 GPU: 257.1 +-0.6
2 GPUs: 249.5 +-11.5
4 GPUs: 447.2 +-56.4

Commit message: Bugfix Keras test exclusion logic

Signed-off-by: Alex Sergeev <alsrgv@users.noreply.github.com>

Patch to review:
diff --git a/.buildkite/gen-pipeline.sh b/.buildkite/gen-pipeline.sh
index c0fbddc3e6..af85d3ac5e 100755
--- a/.buildkite/gen-pipeline.sh
+++ b/.buildkite/gen-pipeline.sh
@@ -17,6 +17,9 @@ tests=( \
        test-cpu-openmpi-py2_7-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
        test-cpu-openmpi-py3_5-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
        test-cpu-openmpi-py3_6-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
+       test-cpu-openmpi-py2_7-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
+       test-cpu-openmpi-py3_5-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
+       test-cpu-openmpi-py3_6-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
        test-cpu-openmpi-py2_7-tfhead-kerashead-torchhead-mxnethead-pyspark2_4_0 \
        test-cpu-openmpi-py3_6-tfhead-kerashead-torchhead-mxnethead-pyspark2_4_0 \
        test-cpu-mpich-py3_6-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
@@ -25,6 +28,8 @@ tests=( \
        test-gpu-openmpi-py3_5-tf1_6_0-keras2_1_2-torch0_4_1-mxnet1_4_1-pyspark2_3_2 \
        test-gpu-openmpi-py2_7-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
        test-gpu-openmpi-py3_6-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
+       test-gpu-openmpi-py2_7-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
+       test-gpu-openmpi-py3_6-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
        test-gpu-openmpi-py2_7-tfhead-kerashead-torchhead-mxnethead-pyspark2_4_0 \
        test-gpu-openmpi-py3_6-tfhead-kerashead-torchhead-mxnethead-pyspark2_4_0 \
        test-gpu-mpich-py3_6-tf1_14_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0 \
@@ -95,24 +100,33 @@ run_all() {
   local test=$1
   local queue=$2
 
+  local exclude_keras_if_needed=""
+  if [[ ${test} == *"tf2_"* ]]; then
+    # TODO: support for Keras + TF 2.0 and TF-Keras 2.0
+    exclude_keras_if_needed="| sed 's/[a-z_]*keras[a-z_.]*//g'"
+  fi
+
   run_test "${test}" "${queue}" \
     ":pytest: Run PyTests (${test})" \
-    "bash -c \"cd /horovod/test && (echo test_*.py | xargs -n 1 \\\$(cat /mpirun_command) pytest -v --capture=no)\""
+    "bash -c \"cd /horovod/test && (echo test_*.py ${exclude_keras_if_needed} | xargs -n 1 \\\$(cat /mpirun_command) pytest -v --capture=no)\""
 
-  run_test "${test}" "${queue}" \
-    ":muscle: Test TensorFlow MNIST (${test})" \
-    "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/tensorflow_mnist.py\""
+  # Legacy TensorFlow tests
+  if [[ ${test} != *"tf2_"* ]]; then
+    run_test "${test}" "${queue}" \
+      ":muscle: Test TensorFlow MNIST (${test})" \
+      "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/tensorflow_mnist.py\""
+
+    if [[ ${test} != *"tf1_1_0"* && ${test} != *"tf1_6_0"* ]]; then
+      run_test "${test}" "${queue}" \
+        ":muscle: Test TensorFlow Eager MNIST (${test})" \
+        "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/tensorflow_mnist_eager.py\""
+    fi
 
-  if [[ ${test} != *"tf1_1_0"* && ${test} != *"tf1_6_0"* ]]; then
     run_test "${test}" "${queue}" \
-      ":muscle: Test TensorFlow Eager MNIST (${test})" \
-      "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/tensorflow_mnist_eager.py\""
+      ":muscle: Test Keras MNIST (${test})" \
+      "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/keras_mnist_advanced.py\""
   fi
 
-  run_test "${test}" "${queue}" \
-    ":muscle: Test Keras MNIST (${test})" \
-    "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/keras_mnist_advanced.py\""
-
   run_test "${test}" "${queue}" \
     ":muscle: Test PyTorch MNIST (${test})" \
     "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/pytorch_mnist.py\""
@@ -134,6 +148,13 @@ run_all() {
         "horovodrun -np 2 -H localhost:2 python /horovod/examples/tensorflow_mnist.py"
     fi
   fi
+
+  # TensorFlow 2.0 tests
+  if [[ ${test} == *"tf2_"* ]]; then
+    run_test "${test}" "${queue}" \
+      ":muscle: Test TensorFlow 2.0 MNIST (${test})" \
+      "bash -c \"\\\$(cat /mpirun_command) python /horovod/examples/tensorflow2_mnist.py\""
+  fi
 }
 
 build_docs() {
diff --git a/docker-compose.test.yml b/docker-compose.test.yml
index a3a2b95279..b53835de66 100644
--- a/docker-compose.test.yml
+++ b/docker-compose.test.yml
@@ -116,6 +116,43 @@ services:
         TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cpu/torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl
         MXNET_PACKAGE: mxnet==1.4.1
         PYSPARK_PACKAGE: pyspark==2.4.0
+  test-cpu-openmpi-py2_7-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0:
+    extends: test-cpu-base
+    build:
+      args:
+        MPI_KIND: OpenMPI
+        PYTHON_VERSION: 2.7
+        TENSORFLOW_PACKAGE: tensorflow==2.0.0b1
+        KERAS_PACKAGE: keras==2.2.4
+        PYTORCH_PACKAGE: torch==1.1.0
+        TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cpu/torchvision-0.3.0-cp27-cp27mu-linux_x86_64.whl
+        MXNET_PACKAGE: mxnet==1.4.1
+        PYSPARK_PACKAGE: pyspark==2.4.0
+  test-cpu-openmpi-py3_5-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0:
+    extends: test-cpu-base
+    build:
+      args:
+        MPI_KIND: OpenMPI
+        PYTHON_VERSION: 3.5
+        TENSORFLOW_PACKAGE: tensorflow==2.0.0b1
+        KERAS_PACKAGE: keras==2.2.4
+        PYTORCH_PACKAGE: torch==1.1.0
+        TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cpu/torchvision-0.3.0-cp35-cp35m-linux_x86_64.whl
+        MXNET_PACKAGE: mxnet==1.4.1
+        PYSPARK_PACKAGE: pyspark==2.4.0
+  test-cpu-openmpi-py3_6-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0:
+    extends: test-cpu-base
+    build:
+      args:
+        UBUNTU_VERSION: 18.04
+        MPI_KIND: OpenMPI
+        PYTHON_VERSION: 3.6
+        TENSORFLOW_PACKAGE: tensorflow==2.0.0b1
+        KERAS_PACKAGE: keras==2.2.4
+        PYTORCH_PACKAGE: torch==1.1.0
+        TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cpu/torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl
+        MXNET_PACKAGE: mxnet==1.4.1
+        PYSPARK_PACKAGE: pyspark==2.4.0
   test-cpu-openmpi-py2_7-tfhead-kerashead-torchhead-mxnethead-pyspark2_4_0:
     extends: test-cpu-base
     build:
@@ -233,6 +270,31 @@ services:
         TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cu100/torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl
         MXNET_PACKAGE: mxnet-cu100==1.4.1
         PYSPARK_PACKAGE: pyspark==2.4.0
+  test-gpu-openmpi-py2_7-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0:
+    extends: test-gpu-base
+    build:
+      args:
+        MPI_KIND: OpenMPI
+        PYTHON_VERSION: 2.7
+        TENSORFLOW_PACKAGE: tensorflow-gpu==2.0.0b1
+        KERAS_PACKAGE: keras==2.2.4
+        PYTORCH_PACKAGE: https://download.pytorch.org/whl/cu100/torch-1.1.0-cp27-cp27mu-linux_x86_64.whl
+        TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cu100/torchvision-0.3.0-cp27-cp27mu-linux_x86_64.whl
+        MXNET_PACKAGE: mxnet-cu100==1.4.1
+        PYSPARK_PACKAGE: pyspark==2.4.0
+  test-gpu-openmpi-py3_6-tf2_0_0-keras2_2_4-torch1_1_0-mxnet1_4_1-pyspark2_4_0:
+    extends: test-gpu-base
+    build:
+      args:
+        CUDA_DOCKER_VERSION: 10.0-devel-ubuntu18.04
+        MPI_KIND: OpenMPI
+        PYTHON_VERSION: 3.6
+        TENSORFLOW_PACKAGE: tensorflow-gpu==2.0.0b1
+        KERAS_PACKAGE: keras==2.2.4
+        PYTORCH_PACKAGE: https://download.pytorch.org/whl/cu100/torch-1.1.0-cp36-cp36m-linux_x86_64.whl
+        TORCHVISION_PACKAGE: https://download.pytorch.org/whl/cu100/torchvision-0.3.0-cp36-cp36m-linux_x86_64.whl
+        MXNET_PACKAGE: mxnet-cu100==1.4.1
+        PYSPARK_PACKAGE: pyspark==2.4.0
   test-gpu-openmpi-py2_7-tfhead-kerashead-torchhead-mxnethead-pyspark2_4_0:
     extends: test-gpu-base
     build:
diff --git a/docs/conf.py b/docs/conf.py
index b7fc507f52..1a9f3bde6a 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -64,7 +64,7 @@
     'special-members': '__init__',
     'imported-members': None,
     'undoc-members': None,
-    'exclude-members': 'contextmanager',
+    'exclude-members': 'contextmanager, LooseVersion, tf, keras, torch, mx, pyspark',
 }
 
 
diff --git a/docs/mocks.py b/docs/mocks.py
index a850e95313..ecf99e5fc2 100644
--- a/docs/mocks.py
+++ b/docs/mocks.py
@@ -45,10 +45,13 @@
 
 MOCK_TREE = {
     'tensorflow': {
-        '__version__': '1.12.0',
+        '__version__': '1.14.0',
         'train': {
             'Optimizer': MagicMock,
-            'SessionRunMock': MagicMock,
+            'SessionRunHook': MagicMock,
+        },
+        'estimator': {
+            'SessionRunHook': MagicMock,
         },
         'keras': {
             'callbacks': {
diff --git a/examples/tensorflow2_mnist.py b/examples/tensorflow2_mnist.py
new file mode 100644
index 0000000000..a7552e17a3
--- /dev/null
+++ b/examples/tensorflow2_mnist.py
@@ -0,0 +1,86 @@
+#!/usr/bin/env python
+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import tensorflow as tf
+import horovod.tensorflow as hvd
+
+# Horovod: initialize Horovod.
+hvd.init()
+
+# Horovod: pin GPU to be used to process local rank (one GPU per process)
+gpus = tf.config.experimental.list_physical_devices('GPU')
+for gpu in gpus:
+    tf.config.experimental.set_memory_growth(gpu, True)
+if gpus:
+    tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
+
+mnist_model = tf.keras.Sequential([
+    tf.keras.layers.Conv2D(16, [3, 3], activation='relu'),
+    tf.keras.layers.Conv2D(16, [3, 3], activation='relu'),
+    tf.keras.layers.GlobalAveragePooling2D(),
+    tf.keras.layers.Dense(10)
+])
+loss = tf.losses.SparseCategoricalCrossentropy(from_logits=True)
+
+# Horovod: adjust learning rate based on number of GPUs.
+opt = tf.optimizers.Adam(0.001 * hvd.size())
+
+(mnist_images, mnist_labels), _ = \
+    tf.keras.datasets.mnist.load_data(path='mnist-%d.npz' % hvd.rank())
+
+dataset = tf.data.Dataset.from_tensor_slices(
+    (tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32),
+        tf.cast(mnist_labels, tf.int64))
+)
+dataset = dataset.repeat().shuffle(1000).batch(32)
+
+checkpoint_dir = './checkpoints'
+checkpoint = tf.train.Checkpoint(model=mnist_model, optimizer=opt)
+
+
+@tf.function
+def training_step(images, labels, first_batch):
+    with tf.GradientTape() as tape:
+        logits = mnist_model(images, training=True)
+        loss_value = loss(labels, logits)
+
+    # Horovod: add Horovod Distributed GradientTape.
+    tape = hvd.DistributedGradientTape(tape)
+
+    grads = tape.gradient(loss_value, mnist_model.trainable_variables)
+    opt.apply_gradients(zip(grads, mnist_model.trainable_variables))
+
+    # Horovod: broadcast initial variable states from rank 0 to all other processes.
+    # This is necessary to ensure consistent initialization of all workers when
+    # training is started with random weights or restored from a checkpoint.
+    if first_batch:
+        hvd.broadcast_variables(mnist_model.variables, root_rank=0)
+        hvd.broadcast_variables(opt.variables(), root_rank=0)
+
+    return loss_value
+
+
+# Horovod: adjust number of steps based on number of GPUs.
+for batch, (images, labels) in enumerate(dataset.take(20000 // hvd.size())):
+    loss_value = training_step(images, labels, batch == 0)
+
+    if batch % 10 == 0 and hvd.local_rank() == 0:
+        print('Step #%d\tLoss: %.6f' % (batch, loss_value))
+
+# Horovod: save checkpoints only on worker 0 to prevent other workers from
+# corrupting it.
+if hvd.rank() == 0:
+    checkpoint.save(checkpoint_dir)
diff --git a/examples/tensorflow2_synthetic_benchmark.py b/examples/tensorflow2_synthetic_benchmark.py
new file mode 100644
index 0000000000..e594368dae
--- /dev/null
+++ b/examples/tensorflow2_synthetic_benchmark.py
@@ -0,0 +1,129 @@
+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+from __future__ import absolute_import, division, print_function
+
+import argparse
+import os
+import numpy as np
+import timeit
+
+import tensorflow as tf
+import horovod.tensorflow as hvd
+from tensorflow.keras import applications
+
+# Benchmark settings
+parser = argparse.ArgumentParser(description='TensorFlow Synthetic Benchmark',
+                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter)
+parser.add_argument('--fp16-allreduce', action='store_true', default=False,
+                    help='use fp16 compression during allreduce')
+
+parser.add_argument('--model', type=str, default='ResNet50',
+                    help='model to benchmark')
+parser.add_argument('--batch-size', type=int, default=32,
+                    help='input batch size')
+
+parser.add_argument('--num-warmup-batches', type=int, default=10,
+                    help='number of warm-up batches that don\'t count towards benchmark')
+parser.add_argument('--num-batches-per-iter', type=int, default=10,
+                    help='number of batches per benchmark iteration')
+parser.add_argument('--num-iters', type=int, default=10,
+                    help='number of benchmark iterations')
+
+parser.add_argument('--no-cuda', action='store_true', default=False,
+                    help='disables CUDA training')
+
+args = parser.parse_args()
+args.cuda = not args.no_cuda
+
+# Horovod: initialize Horovod.
+hvd.init()
+
+# Horovod: pin GPU to be used to process local rank (one GPU per process)
+if args.cuda:
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    for gpu in gpus:
+        tf.config.experimental.set_memory_growth(gpu, True)
+    if gpus:
+        tf.config.experimental.set_visible_devices(gpus[hvd.local_rank()], 'GPU')
+else:
+    os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
+
+# Set up standard model.
+model = getattr(applications, args.model)(weights=None)
+opt = tf.optimizers.SGD(0.01)
+
+data = tf.random.uniform([args.batch_size, 224, 224, 3])
+target = tf.random.uniform([args.batch_size, 1], minval=0, maxval=999, dtype=tf.int64)
+
+
+@tf.function
+def benchmark_step(first_batch):
+    # Horovod: (optional) compression algorithm.
+    compression = hvd.Compression.fp16 if args.fp16_allreduce else hvd.Compression.none
+
+    # Horovod: use DistributedGradientTape
+    with tf.GradientTape() as tape:
+        logits = model(data, training=True)
+        loss = tf.losses.categorical_crossentropy(target, logits)
+
+    # Horovod: add Horovod Distributed GradientTape.
+    tape = hvd.DistributedGradientTape(tape, compression=compression)
+
+    gradients = tape.gradient(loss, model.trainable_variables)
+    opt.apply_gradients(zip(gradients, model.trainable_variables))
+
+    # Horovod: broadcast initial variable states from rank 0 to all other processes.
+    # This is necessary to ensure consistent initialization of all workers when
+    # training is started with random weights or restored from a checkpoint.
+    if first_batch:
+        hvd.broadcast_variables(model.variables, root_rank=0)
+        hvd.broadcast_variables(opt.variables(), root_rank=0)
+
+
+def log(s, nl=True):
+    if hvd.rank() != 0:
+        return
+    print(s, end='\n' if nl else '')
+
+
+log('Model: %s' % args.model)
+log('Batch size: %d' % args.batch_size)
+device = 'GPU' if args.cuda else 'CPU'
+log('Number of %ss: %d' % (device, hvd.size()))
+
+
+with tf.device(device):
+    # Warm-up
+    log('Running warmup...')
+    benchmark_step(first_batch=True)
+    timeit.timeit(lambda: benchmark_step(first_batch=False),
+                  number=args.num_warmup_batches)
+
+    # Benchmark
+    log('Running benchmark...')
+    img_secs = []
+    for x in range(args.num_iters):
+        time = timeit.timeit(lambda: benchmark_step(first_batch=False),
+                             number=args.num_batches_per_iter)
+        img_sec = args.batch_size * args.num_batches_per_iter / time
+        log('Iter #%d: %.1f img/sec per %s' % (x, img_sec, device))
+        img_secs.append(img_sec)
+
+    # Results
+    img_sec_mean = np.mean(img_secs)
+    img_sec_conf = 1.96 * np.std(img_secs)
+    log('Img/sec per %s: %.1f +-%.1f' % (device, img_sec_mean, img_sec_conf))
+    log('Total img/sec on %d %s(s): %.1f +-%.1f' %
+        (hvd.size(), device, hvd.size() * img_sec_mean, hvd.size() * img_sec_conf))
diff --git a/examples/tensorflow_mnist.py b/examples/tensorflow_mnist.py
index cd271040d6..3d83c6d84c 100644
--- a/examples/tensorflow_mnist.py
+++ b/examples/tensorflow_mnist.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright 2017 Uber Technologies, Inc. All Rights Reserved.
+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -113,7 +113,7 @@ def main(_):
     predict, loss = conv_model(image, label, tf.estimator.ModeKeys.TRAIN)
 
     # Horovod: adjust learning rate based on number of GPUs.
-    opt = tf.train.RMSPropOptimizer(0.001 * hvd.size())
+    opt = tf.train.AdamOptimizer(0.001 * hvd.size())
 
     # Horovod: add Horovod Distributed Optimizer.
     opt = hvd.DistributedOptimizer(opt)
diff --git a/examples/tensorflow_mnist_eager.py b/examples/tensorflow_mnist_eager.py
index 66d47b0a6d..605f274d33 100644
--- a/examples/tensorflow_mnist_eager.py
+++ b/examples/tensorflow_mnist_eager.py
@@ -1,5 +1,5 @@
 #!/usr/bin/env python
-# Copyright 2018 Uber Technologies, Inc. All Rights Reserved.
+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -36,7 +36,7 @@ def main(_):
     ])
 
     # Horovod: adjust learning rate based on number of GPUs.
-    opt = tf.train.RMSPropOptimizer(0.001 * hvd.size())
+    opt = tf.train.AdamOptimizer(0.001 * hvd.size())
 
     (mnist_images, mnist_labels), _ = \
         tf.keras.datasets.mnist.load_data(path='mnist-%d.npz' % hvd.rank())
@@ -45,7 +45,7 @@ def main(_):
         (tf.cast(mnist_images[..., tf.newaxis] / 255.0, tf.float32),
          tf.cast(mnist_labels, tf.int64))
     )
-    dataset = dataset.shuffle(1000).batch(32)
+    dataset = dataset.repeat().shuffle(1000).batch(32)
 
     checkpoint_dir = './checkpoints'
     step_counter = tf.train.get_or_create_global_step()
@@ -59,12 +59,6 @@ def main(_):
             logits = mnist_model(images, training=True)
             loss_value = tf.losses.sparse_softmax_cross_entropy(labels, logits)
 
-        # Horovod: broadcast initial variable states from rank 0 to all other processes.
-        # This is necessary to ensure consistent initialization of all workers when
-        # training is started with random weights or restored from a checkpoint.
-        if batch == 0:
-            hvd.broadcast_variables(mnist_model.variables, root_rank=0)
-
         # Horovod: add Horovod Distributed GradientTape.
         tape = hvd.DistributedGradientTape(tape)
 
@@ -72,6 +66,13 @@ def main(_):
         opt.apply_gradients(zip(grads, mnist_model.variables),
                             global_step=tf.train.get_or_create_global_step())
 
+        # Horovod: broadcast initial variable states from rank 0 to all other processes.
+        # This is necessary to ensure consistent initialization of all workers when
+        # training is started with random weights or restored from a checkpoint.
+        if batch == 0:
+            hvd.broadcast_variables(mnist_model.variables, root_rank=0)
+            hvd.broadcast_variables(opt.variables(), root_rank=0)
+
         if batch % 10 == 0 and hvd.local_rank() == 0:
             print('Step #%d\tLoss: %.6f' % (batch, loss_value))
 
diff --git a/horovod/tensorflow/__init__.py b/horovod/tensorflow/__init__.py
index a9b92ead1b..5c14eedbb7 100644
--- a/horovod/tensorflow/__init__.py
+++ b/horovod/tensorflow/__init__.py
@@ -29,7 +29,7 @@
 from horovod.tensorflow.mpi_ops import init, shutdown
 from horovod.tensorflow.mpi_ops import size, local_size, rank, local_rank
 from horovod.tensorflow.mpi_ops import mpi_threads_supported
-from horovod.tensorflow.util import _executing_eagerly
+from horovod.tensorflow.util import _executing_eagerly, _make_subgraph, _cache
 
 import tensorflow as tf
 
@@ -68,7 +68,7 @@ def allreduce(tensor, average=True, device_dense='', device_sparse='',
 
             # To make this operation into an average, divide allgathered values by
             # the Horovod size.
-            new_values = tf.div(values, horovod_size) if average else values
+            new_values = (values / horovod_size) if average else values
         return tf.IndexedSlices(new_values, indices,
                                 dense_shape=tensor.dense_shape)
     else:
@@ -77,19 +77,19 @@ def allreduce(tensor, average=True, device_dense='', device_sparse='',
             tensor_compressed, ctx = compression.compress(tensor)
             summed_tensor_compressed = _allreduce(tensor_compressed)
             summed_tensor = compression.decompress(summed_tensor_compressed, ctx)
-            new_tensor = (tf.div(summed_tensor, horovod_size)
-                          if average else summed_tensor)
+            new_tensor = (summed_tensor / horovod_size) if average else summed_tensor
         return new_tensor
 
 
-def broadcast_global_variables(root_rank):
-    """Broadcasts all global variables from root rank to all other processes.
+@_cache
+def _make_broadcast_group_fn():
+    def broadcast_group(variables, root_rank):
+        return [var.assign(broadcast(var, root_rank)) for var in variables]
 
-    Arguments:
-        root_rank: rank of the process from which global variables will be broadcasted
-        to all other processes.
-    """
-    return broadcast_variables(tf.global_variables(), root_rank)
+    if _executing_eagerly():
+        return _make_subgraph(broadcast_group)
+    else:
+        return broadcast_group
 
 
 def broadcast_variables(variables, root_rank):
@@ -100,195 +100,164 @@ def broadcast_variables(variables, root_rank):
         root_rank: rank of the process from which global variables will be broadcasted
                    to all other processes.
     """
-    return tf.group(*[tf.assign(var, broadcast(var, root_rank))
-                      for var in variables])
+    broadcast_group = _make_broadcast_group_fn()
+    return broadcast_group(variables, root_rank)
 
 
-class BroadcastGlobalVariablesHook(tf.train.SessionRunHook):
-    """
-    SessionRunHook that will broadcast all global variables from root rank
-    to all other processes during initialization.
+if hasattr(tf, 'global_variables'):
+    def broadcast_global_variables(root_rank):
+        """Broadcasts all global variables from root rank to all other processes.
 
-    This is necessary to ensure consistent initialization of all workers when
-    training is started with random weights or restored from a checkpoint.
-    """
+        **NOTE:** deprecated in TensorFlow 2.0.
 
-    def __init__(self, root_rank, device=''):
-        """Construct a new BroadcastGlobalVariablesHook that will broadcast all
-        global variables from root rank to all other processes during initialization.
-
-        Args:
-          root_rank:
-            Rank that will send data, other ranks will receive data.
-          device:
-            Device to be used for broadcasting. Uses GPU by default
-            if Horovod was build with HOROVOD_GPU_BROADCAST.
+        Arguments:
+            root_rank: rank of the process from which global variables will be broadcasted
+                       to all other processes.
         """
-        super(BroadcastGlobalVariablesHook, self).__init__()
-        self.root_rank = root_rank
-        self.bcast_op = None
-        self.device = device
-
-    def begin(self):
-        if not self.bcast_op or self.bcast_op.graph != tf.get_default_graph():
-            with tf.device(self.device):
-                self.bcast_op = broadcast_global_variables(self.root_rank)
-
-    def after_create_session(self, session, coord):
-        session.run(self.bcast_op)
-
+        return broadcast_variables(tf.global_variables(), root_rank)
 
-class DistributedOptimizer(tf.train.Optimizer):
-    """An optimizer that wraps another tf.Optimizer, using an allreduce to
-    average gradient values before applying gradients to model weights."""
 
-    def __init__(self, optimizer, name=None, use_locking=False, device_dense='',
-                 device_sparse='', compression=Compression.none,
-                 sparse_as_dense=False):
-        """Construct a new DistributedOptimizer, which uses another optimizer
-        under the hood for computing single-process gradient values and
-        applying gradient updates after the gradient values have been averaged
-        across all the Horovod ranks.
+if hasattr(tf, 'train') and hasattr(tf.train, 'SessionRunHook'):
+    if hasattr(tf, 'estimator') and hasattr(tf.estimator, 'SessionRunHook'):
+        _SessionRunHook = tf.estimator.SessionRunHook
+    else:
+        _SessionRunHook = tf.train.SessionRunHook
 
-        Args:
-          optimizer:
-            Optimizer to use for computing gradients and applying updates.
-          name:
-            Optional name prefix for the operations created when applying
-            gradients. Defaults to "Distributed" followed by the provided
-            optimizer type.
-          use_locking:
-            Whether to use locking when updating variables.
-            See Optimizer.__init__ for more info.
-          device_dense:
-            Device to be used for dense tensors. Uses GPU by default
-            if Horovod was build with HOROVOD_GPU_ALLREDUCE.
-          device_sparse:
-            Device to be used for sparse tensors. Uses GPU by default
-            if Horovod was build with HOROVOD_GPU_ALLGATHER.
-          compression:
-            Compression algorithm used during allreduce to reduce the amount
-            of data sent during the each parameter update step.  Defaults to
-            not using compression.
-          sparse_as_dense:
-            Treat all sparse gradients as dense tensors.  This can help improve
-            performance and memory utilization if the original sparse gradient
-            has high density.  Defaults to false.
+    class BroadcastGlobalVariablesHook(_SessionRunHook):
         """
-        if name is None:
-            name = "Distributed{}".format(type(optimizer).__name__)
-
-        self._optimizer = optimizer
-        self._device_dense = device_dense
-        self._device_sparse = device_sparse
-        self._compression = compression
-        self._sparse_as_dense = sparse_as_dense
-
-        def allreduce_grads(grads):
-            with tf.name_scope(self._name + "_Allreduce"):
-                if self._sparse_as_dense:
-                    grads = [tf.convert_to_tensor(grad)
-                             if grad is not None and isinstance(grad, tf.IndexedSlices)
-                             else grad for grad in grads]
-
-                return [allreduce(grad,
-                                  device_dense=self._device_dense,
-                                  device_sparse=self._device_sparse,
-                                  compression=self._compression)
-                        if grad is not None else grad
-                        for grad in grads]
-
-        if _executing_eagerly():
-            self._allreduce_grads = tf.contrib.eager.defun(allreduce_grads)
-        else:
-            self._allreduce_grads = allreduce_grads
+        SessionRunHook that will broadcast all global variables from root rank
+        to all other processes during initialization.
 
-        super(DistributedOptimizer, self).__init__(
-            name=name, use_locking=use_locking)
+        This is necessary to ensure consistent initialization of all workers when
+        training is started with random weights or restored from a checkpoint.
 
-    def compute_gradients(self, *args, **kwargs):
-        """Compute gradients of all trainable variables.
+        **NOTE:** deprecated in TensorFlow 2.0.
+        """
 
-        See Optimizer.compute_gradients() for more info.
+        def __init__(self, root_rank, device=''):
+            """Construct a new BroadcastGlobalVariablesHook that will broadcast all
+            global variables from root rank to all other processes during initialization.
+
+            Args:
+              root_rank:
+                Rank that will send data, other ranks will receive data.
+              device:
+                Device to be used for broadcasting. Uses GPU by default
+                if Horovod was build with HOROVOD_GPU_BROADCAST.
+            """
+            super(BroadcastGlobalVariablesHook, self).__init__()
+            self.root_rank = root_rank
+            self.bcast_op = None
+            self.device = device
+
+        def begin(self):
+            if not self.bcast_op or self.bcast_op.graph != tf.get_default_graph():
+                with tf.device(self.device):
+                    self.bcast_op = broadcast_global_variables(self.root_rank)
+
+        def after_create_session(self, session, coord):
+            session.run(self.bcast_op)
+
+
+@_cache
+def _make_allreduce_grads_fn(name, device_dense, device_sparse,
+                             compression, sparse_as_dense):
+    def allreduce_grads(grads):
+        with tf.name_scope(name + "_Allreduce"):
+            if sparse_as_dense:
+                grads = [tf.convert_to_tensor(grad)
+                         if grad is not None and isinstance(grad, tf.IndexedSlices)
+                         else grad for grad in grads]
+
+            return [allreduce(grad,
+                              device_dense=device_dense,
+                              device_sparse=device_sparse,
+                              compression=compression)
+                    if grad is not None else grad
+                    for grad in grads]
+
+    if _executing_eagerly():
+        return _make_subgraph(allreduce_grads)
+    else:
+        return allreduce_grads
 
-        In DistributedOptimizer, compute_gradients() is overriden to also
-        allreduce the gradients before returning them.
-        """
-        gradients = self._optimizer.compute_gradients(*args, **kwargs)
-        if size() > 1:
-            grads, vars = zip(*gradients)
-            avg_grads = self._allreduce_grads(grads)
-            return list(zip(avg_grads, vars))
-        else:
-            return gradients
 
-    def apply_gradients(self, *args, **kwargs):
-        """Calls this same method on the underlying optimizer."""
-        return self._optimizer.apply_gradients(*args, **kwargs)
+if hasattr(tf, 'compat') and hasattr(tf.compat, 'v1') and \
+        hasattr(tf.compat.v1, 'train') and hasattr(tf.compat.v1.train, 'Optimizer'):
+    _LegacyOptimizer = tf.compat.v1.train.Optimizer
+elif hasattr(tf, 'train') and hasattr(tf.train, 'Optimizer'):
+    _LegacyOptimizer = tf.train.Optimizer
+else:
+    _LegacyOptimizer = None
 
-    def get_slot(self, *args, **kwargs):
-        """Calls this same method on the underlying optimizer."""
-        return self._optimizer.get_slot(*args, **kwargs)
 
-    def get_slot_names(self, *args, **kwargs):
-        """Calls this same method on the underlying optimizer."""
-        return self._optimizer.get_slot_names(*args, **kwargs)
+if _LegacyOptimizer:
+    class _DistributedOptimizer(_LegacyOptimizer):
+        """An optimizer that wraps another tf.Optimizer, using an allreduce to
+        average gradient values before applying gradients to model weights."""
 
-    def variables(self, *args, **kwargs):
-        """Calls this same method on the underlying optimizer."""
-        return self._optimizer.variables(*args, **kwargs)
+        def __init__(self, optimizer, name=None, use_locking=False, device_dense='',
+                    device_sparse='', compression=Compression.none,
+                    sparse_as_dense=False):
+            if name is None:
+                name = "Distributed{}".format(type(optimizer).__name__)
+            super(_DistributedOptimizer, self).__init__(name=name, use_locking=use_locking)
 
+            self._optimizer = optimizer
+            self._allreduce_grads = _make_allreduce_grads_fn(
+                name, device_dense, device_sparse, compression, sparse_as_dense)
 
-if hasattr(tf, 'GradientTape'):
-    class _DistributedGradientTape(tf.GradientTape):
+        def compute_gradients(self, *args, **kwargs):
+            """Compute gradients of all trainable variables.
 
-        def __init__(self, tape, device_dense, device_sparse,
-                     compression, sparse_as_dense, persistent=False, watch_accessed_variables=True):
-            if hasattr(tape, '_watch_accessed_variables'):
-                super(self.__class__, self).__init__(persistent, watch_accessed_variables)
-            else:
-                super(self.__class__, self).__init__(persistent)
-            self._tape = tape
-            self._persistent = persistent
-            self._watch_accessed_variables = watch_accessed_variables
-            self._name = "Distributed"
-            self._device_dense = device_dense
-            self._device_sparse = device_sparse
-            self._compression = compression
-            self._sparse_as_dense = sparse_as_dense
-
-            def allreduce_grads(grads):
-                with tf.name_scope(self._name + "_Allreduce"):
-                    if self._sparse_as_dense:
-                        grads = [tf.convert_to_tensor(grad)
-                                 if grad is not None and isinstance(grad, tf.IndexedSlices)
-                                 else grad for grad in grads]
-                    return [allreduce(grad,
-                                      device_dense=self._device_dense,
-                                      device_sparse=self._device_sparse,
-                                      compression=self._compression)
-                            if grad is not None else grad
-                            for grad in grads]
-
-            self._allreduce_grads = tf.contrib.eager.defun(allreduce_grads)
+            See Optimizer.compute_gradients() for more info.
 
-        def gradient(self, target, sources, output_gradients=None):
-            gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
+            In DistributedOptimizer, compute_gradients() is overriden to also
+            allreduce the gradients before returning them.
+            """
+            gradients = self._optimizer.compute_gradients(*args, **kwargs)
             if size() > 1:
-                avg_grads = self._allreduce_grads(gradients)
-                return avg_grads
+                grads, vars = zip(*gradients)
+                avg_grads = self._allreduce_grads(grads)
+                return list(zip(avg_grads, vars))
             else:
                 return gradients
 
+        def apply_gradients(self, *args, **kwargs):
+            """Calls this same method on the underlying optimizer."""
+            return self._optimizer.apply_gradients(*args, **kwargs)
 
-def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
-                            compression=Compression.none, sparse_as_dense=False):
-    """An tape that wraps another tf.GradientTape, using an allreduce to
-    average gradient values before applying gradients to model weights.
+        def get_slot(self, *args, **kwargs):
+            """Calls this same method on the underlying optimizer."""
+            return self._optimizer.get_slot(*args, **kwargs)
+
+        def get_slot_names(self, *args, **kwargs):
+            """Calls this same method on the underlying optimizer."""
+            return self._optimizer.get_slot_names(*args, **kwargs)
+
+        def variables(self, *args, **kwargs):
+            """Calls this same method on the underlying optimizer."""
+            return self._optimizer.variables(*args, **kwargs)
+
+
+def DistributedOptimizer(optimizer, name=None, use_locking=False, device_dense='',
+                         device_sparse='', compression=Compression.none,
+                         sparse_as_dense=False):
+    """Construct a new DistributedOptimizer, which uses another optimizer
+    under the hood for computing single-process gradient values and
+    applying gradient updates after the gradient values have been averaged
+    across all the Horovod ranks.
 
     Args:
-      gradtape:
-        GradientTape to use for computing gradients and applying updates.
+      optimizer:
+        Optimizer to use for computing gradients and applying updates.
+      name:
+        Optional name prefix for the operations created when applying
+        gradients. Defaults to "Distributed" followed by the provided
+        optimizer type.
+      use_locking:
+        Whether to use locking when updating variables.
+        See Optimizer.__init__ for more info.
       device_dense:
         Device to be used for dense tensors. Uses GPU by default
         if Horovod was build with HOROVOD_GPU_ALLREDUCE.
@@ -304,13 +273,69 @@ def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
         performance and memory utilization if the original sparse gradient
         has high density.  Defaults to false.
     """
-    cls = type(gradtape.__class__.__name__, (gradtape.__class__,),
-               dict(_DistributedGradientTape.__dict__))
-    if hasattr(gradtape, '_watch_accessed_variables'):
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent, gradtape._watch_accessed_variables)
+    if isinstance(optimizer, _LegacyOptimizer):
+        return _DistributedOptimizer(optimizer, name, use_locking, device_dense,
+                                     device_sparse, compression, sparse_as_dense)
+    elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
+        import horovod.tensorflow.keras as hvd_k
+        return hvd_k.DistributedOptimizer(optimizer, name, device_dense, device_sparse,
+                                          compression, sparse_as_dense)
     else:
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent)
+        raise ValueError('Provided optimizer doesn\'t inherit from either legacy '
+                         'TensorFlow or Keras optimizer: %s' % optimizer)
+
+
+if hasattr(tf, 'GradientTape'):
+    class _DistributedGradientTape(tf.GradientTape):
+        def __init__(self, tape, device_dense, device_sparse, compression, sparse_as_dense,
+                     persistent=False, watch_accessed_variables=True):
+            if hasattr(tape, '_watch_accessed_variables'):
+                super(self.__class__, self).__init__(persistent, watch_accessed_variables)
+            else:
+                super(self.__class__, self).__init__(persistent)
+
+            self._tape = tape
+            self._allreduce_grads = _make_allreduce_grads_fn(
+                'DistributedGradientTape', device_dense, device_sparse, compression,
+                sparse_as_dense)
+
+        def gradient(self, target, sources, output_gradients=None):
+            gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
+            if size() > 1:
+                return self._allreduce_grads(gradients)
+            else:
+                return gradients
+
+
+    def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
+                                compression=Compression.none, sparse_as_dense=False):
+        """An tape that wraps another tf.GradientTape, using an allreduce to
+        average gradient values before applying gradients to model weights.
+
+        Args:
+          gradtape:
+            GradientTape to use for computing gradients and applying updates.
+          device_dense:
+            Device to be used for dense tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLREDUCE.
+          device_sparse:
+            Device to be used for sparse tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLGATHER.
+          compression:
+            Compression algorithm used during allreduce to reduce the amount
+            of data sent during the each parameter update step.  Defaults to
+            not using compression.
+          sparse_as_dense:
+            Treat all sparse gradients as dense tensors.  This can help improve
+            performance and memory utilization if the original sparse gradient
+            has high density.  Defaults to false.
+        """
+        cls = type(gradtape.__class__.__name__, (gradtape.__class__,),
+                   dict(_DistributedGradientTape.__dict__))
+        if hasattr(gradtape, '_watch_accessed_variables'):
+            return cls(gradtape._tape, device_dense, device_sparse, compression,
+                       sparse_as_dense, gradtape._persistent,
+                       gradtape._watch_accessed_variables)
+        else:
+            return cls(gradtape._tape, device_dense, device_sparse, compression,
+                       sparse_as_dense, gradtape._persistent)
diff --git a/horovod/tensorflow/util.py b/horovod/tensorflow/util.py
index 127403db91..96e23ab0a9 100644
--- a/horovod/tensorflow/util.py
+++ b/horovod/tensorflow/util.py
@@ -1,4 +1,4 @@
-# Copyright 2018 Uber Technologies, Inc. All Rights Reserved.
+# Copyright 2019 Uber Technologies, Inc. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,3 +27,26 @@
 def _executing_eagerly():
     """Returns true if eager execution is supported and enabled."""
     return _has_eager and context.in_eager_mode()
+
+
+def _make_subgraph(f):
+    if hasattr(tf, 'function'):
+        return tf.function(f)
+    else:
+        return tf.contrib.eager.defun(f)
+
+
+def _cache(f):
+    cache = dict()
+
+    def wrapper(*args):
+        key = (args, _executing_eagerly())
+
+        if key in cache:
+            return cache[key]
+        else:
+            retval = f(*args)
+            cache[key] = retval
+            return retval
+
+    return wrapper
diff --git a/test/test_tensorflow.py b/test/test_tensorflow.py
index 87e6e93cea..cc60bc684e 100644
--- a/test/test_tensorflow.py
+++ b/test/test_tensorflow.py
@@ -33,13 +33,20 @@
 
 from common import mpi_env_rank_and_size
 
-config = tf.ConfigProto()
-config.gpu_options.allow_growth = True
-
-if _has_eager:
-    # Specifies the config to use with eager execution. Does not preclude
-    # tests from running in the graph mode.
-    tf.enable_eager_execution(config=config)
+if hasattr(tf, 'ConfigProto'):
+    config = tf.ConfigProto()
+    config.gpu_options.allow_growth = True
+
+if hasattr(tf, 'config') and hasattr(tf.config, 'experimental') \
+        and hasattr(tf.config.experimental, 'set_memory_growth'):
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    for gpu in gpus:
+        tf.config.experimental.set_memory_growth(gpu, True)
+else:
+    if _has_eager:
+        # Specifies the config to use with eager execution. Does not preclude
+        # tests from running in the graph mode.
+        tf.enable_eager_execution(config=config)
 
 # MLSL supports only byte, float and double data types
 mlsl_supported_types = set([tf.float32, tf.float64])
@@ -54,7 +61,10 @@ def __init__(self, *args, **kwargs):
         super(MPITests, self).__init__(*args, **kwargs)
         warnings.simplefilter('module')
         if _has_eager:
-            self.tfe = tf.contrib.eager
+            if hasattr(tf, 'contrib') and hasattr(tf.contrib, 'eager'):
+                self.tfe = tf.contrib.eager
+            else:
+                self.tfe = tf
 
     def evaluate(self, tensors):
         if _executing_eagerly():
@@ -66,6 +76,14 @@ def evaluate(self, tensors):
         else:
             return sess.run(tensors)
 
+    def random_uniform(self, *args, **kwargs):
+        if hasattr(tf, 'random') and hasattr(tf.random, 'set_seed'):
+            tf.random.set_seed(1234)
+            return tf.random.uniform(*args, **kwargs)
+        else:
+            tf.set_random_seed(1234)
+            return tf.random_uniform(*args, **kwargs)
+
     def filter_supported_types(self, types):
         if 'MLSL_ROOT' in os.environ:
            types = [t for t in types if t in mlsl_supported_types]
@@ -93,8 +111,7 @@ def test_horovod_allreduce_cpu(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             with tf.device("/cpu:0"):
-                tf.set_random_seed(1234)
-                tensor = tf.random_uniform(
+                tensor = self.random_uniform(
                     [17] * dim, -100, 100, dtype=dtype)
                 summed = hvd.allreduce(tensor, average=False)
             multiplied = tensor * size
@@ -125,8 +142,7 @@ def test_horovod_allreduce_cpu_fused(self):
         tests = []
         for dtype, dim in itertools.product(dtypes, dims):
             with tf.device("/cpu:0"):
-                tf.set_random_seed(1234)
-                tensor = tf.random_uniform(
+                tensor = self.random_uniform(
                     [17] * dim, -100, 100, dtype=dtype)
                 summed = hvd.allreduce(tensor, average=False)
             multiplied = tensor * size
@@ -166,8 +182,7 @@ def test_horovod_allreduce_gpu(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             with tf.device("/gpu:%d" % local_rank):
-                tf.set_random_seed(1234)
-                tensor = tf.random_uniform(
+                tensor = self.random_uniform(
                     [17] * dim, -100, 100, dtype=dtype)
                 summed = hvd.allreduce(tensor, average=False)
             multiplied = tensor * size
@@ -211,8 +226,7 @@ def test_horovod_allreduce_gpu_fused(self):
         tests = []
         for dtype, dim in itertools.product(dtypes, dims):
             with tf.device("/gpu:%d" % local_rank):
-                tf.set_random_seed(1234)
-                tensor = tf.random_uniform(
+                tensor = self.random_uniform(
                     [17] * dim, -100, 100, dtype=dtype)
                 summed = hvd.allreduce(tensor, average=False)
             multiplied = tensor * size
@@ -259,8 +273,7 @@ def test_horovod_allreduce_multi_gpu(self):
         for dtype, dim in itertools.product(dtypes, dims):
             iter += 1
             with tf.device("/gpu:%d" % gpu_ids[(iter + local_rank) % 2]):
-                tf.set_random_seed(1234)
-                tensor = tf.random_uniform(
+                tensor = self.random_uniform(
                     [17] * dim, -100, 100, dtype=dtype)
                 summed = hvd.allreduce(tensor, average=False)
             multiplied = tensor * size
@@ -293,19 +306,17 @@ def test_horovod_allreduce_error(self):
             return
 
         # Same rank, different dimension
-        tf.set_random_seed(1234)
         dims = [17 + rank] * 3
-        tensor = tf.random_uniform(dims, -1.0, 1.0)
+        tensor = self.random_uniform(dims, -1.0, 1.0)
         with self.assertRaises(tf.errors.FailedPreconditionError):
             self.evaluate(hvd.allreduce(tensor))
 
         # Same number of elements, different rank
-        tf.set_random_seed(1234)
         if rank == 0:
             dims = [17, 23 * 57]
         else:
             dims = [17, 23, 57]
-        tensor = tf.random_uniform(dims, -1.0, 1.0)
+        tensor = self.random_uniform(dims, -1.0, 1.0)
         with self.assertRaises(tf.errors.FailedPreconditionError):
             self.evaluate(hvd.allreduce(tensor))
 
@@ -365,14 +376,13 @@ def test_horovod_allreduce_grad_cpu(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             with tf.device("/cpu:0"):
-                tf.set_random_seed(1234)
                 if _executing_eagerly():
-                    tensor = self.tfe.Variable(tf.random_uniform(
+                    tensor = self.tfe.Variable(self.random_uniform(
                         [5] * dim, -100, 100, dtype=dtype))
                     with tf.GradientTape() as tape:
                         summed = hvd.allreduce(tensor, average=False)
                 else:
-                    tensor = tf.random_uniform(
+                    tensor = self.random_uniform(
                         [5] * dim, -100, 100, dtype=dtype)
                     summed = hvd.allreduce(tensor, average=False)
 
@@ -409,14 +419,13 @@ def test_horovod_allreduce_grad_gpu(self):
         dims = [1, 2, 3]
         for dtype, dim in itertools.product(dtypes, dims):
             with tf.device("/gpu:%d" % local_rank):
-                tf.set_random_seed(1234)
                 if _executing_eagerly():
                     tensor = self.tfe.Variable(
-                        tf.random_uniform([5] * dim, -100, 100, dtype=dtype))
+                        self.random_uniform([5] * dim, -100, 100, dtype=dtype))
                     with tf.GradientTape() as tape:
                         summed = hvd.allreduce(tensor, average=False)
                 else:
-                    tensor = tf.random_uniform([5] * dim, -100, 100, dtype=dtype)
+                    tensor = self.random_uniform([5] * dim, -100, 100, dtype=dtype)
                     summed = hvd.allreduce(tensor, average=False)
 
                 grad_ys = tf.ones([5] * dim)


Code review comments:
diff_hunk: @@ -304,13 +273,69 @@ def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
         performance and memory utilization if the original sparse gradient
         has high density.  Defaults to false.
     """
-    cls = type(gradtape.__class__.__name__, (gradtape.__class__,),
-               dict(_DistributedGradientTape.__dict__))
-    if hasattr(gradtape, '_watch_accessed_variables'):
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent, gradtape._watch_accessed_variables)
+    if isinstance(optimizer, _LegacyOptimizer):
+        return _DistributedOptimizer(optimizer, name, use_locking, device_dense,
+                                     device_sparse, compression, sparse_as_dense)
+    elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
+        import horovod.tensorflow.keras as hvd_k
+        return hvd_k.DistributedOptimizer(optimizer, name, device_dense, device_sparse,
+                                          compression, sparse_as_dense)
     else:
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent)
+        raise ValueError('Provided optimizer doesn\'t inherit from either legacy '
+                         'TensorFlow or Keras optimizer: %s' % optimizer)
+
+
+if hasattr(tf, 'GradientTape'):
+    class _DistributedGradientTape(tf.GradientTape):
+        def __init__(self, tape, device_dense, device_sparse, compression, sparse_as_dense,
+                     persistent=False, watch_accessed_variables=True):
+            if hasattr(tape, '_watch_accessed_variables'):
+                super(self.__class__, self).__init__(persistent, watch_accessed_variables)
+            else:
+                super(self.__class__, self).__init__(persistent)
+
+            self._tape = tape
+            self._allreduce_grads = _make_allreduce_grads_fn(
+                'DistributedGradientTape', device_dense, device_sparse, compression,
+                sparse_as_dense)
+
+        def gradient(self, target, sources, output_gradients=None):
+            gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
+            if size() > 1:
+                return self._allreduce_grads(gradients)
+            else:
+                return gradients
+
+
+    def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
+                                compression=Compression.none, sparse_as_dense=False):
+        """An tape that wraps another tf.GradientTape, using an allreduce to
path: horovod/tensorflow/__init__.py
text: @user1:
"An" -> "A"

diff_hunk: @@ -304,13 +273,69 @@ def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
         performance and memory utilization if the original sparse gradient
         has high density.  Defaults to false.
     """
-    cls = type(gradtape.__class__.__name__, (gradtape.__class__,),
-               dict(_DistributedGradientTape.__dict__))
-    if hasattr(gradtape, '_watch_accessed_variables'):
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent, gradtape._watch_accessed_variables)
+    if isinstance(optimizer, _LegacyOptimizer):
+        return _DistributedOptimizer(optimizer, name, use_locking, device_dense,
+                                     device_sparse, compression, sparse_as_dense)
+    elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
+        import horovod.tensorflow.keras as hvd_k
+        return hvd_k.DistributedOptimizer(optimizer, name, device_dense, device_sparse,
+                                          compression, sparse_as_dense)
     else:
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent)
+        raise ValueError('Provided optimizer doesn\'t inherit from either legacy '
+                         'TensorFlow or Keras optimizer: %s' % optimizer)
+
+
+if hasattr(tf, 'GradientTape'):
+    class _DistributedGradientTape(tf.GradientTape):
+        def __init__(self, tape, device_dense, device_sparse, compression, sparse_as_dense,
+                     persistent=False, watch_accessed_variables=True):
+            if hasattr(tape, '_watch_accessed_variables'):
+                super(self.__class__, self).__init__(persistent, watch_accessed_variables)
+            else:
+                super(self.__class__, self).__init__(persistent)
+
+            self._tape = tape
+            self._allreduce_grads = _make_allreduce_grads_fn(
+                'DistributedGradientTape', device_dense, device_sparse, compression,
+                sparse_as_dense)
+
+        def gradient(self, target, sources, output_gradients=None):
+            gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
+            if size() > 1:
+                return self._allreduce_grads(gradients)
+            else:
+                return gradients
+
+
+    def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
+                                compression=Compression.none, sparse_as_dense=False):
+        """An tape that wraps another tf.GradientTape, using an allreduce to
+        average gradient values before applying gradients to model weights.
+
+        Args:
+          gradtape:
+            GradientTape to use for computing gradients and applying updates.
+          device_dense:
+            Device to be used for dense tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLREDUCE.
+          device_sparse:
+            Device to be used for sparse tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLGATHER.
+          compression:
+            Compression algorithm used during allreduce to reduce the amount
+            of data sent during the each parameter update step.  Defaults to
path: horovod/tensorflow/__init__.py
text: @user1:
"the each" -> "each"

diff_hunk: @@ -304,13 +273,69 @@ def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
         performance and memory utilization if the original sparse gradient
         has high density.  Defaults to false.
     """
-    cls = type(gradtape.__class__.__name__, (gradtape.__class__,),
-               dict(_DistributedGradientTape.__dict__))
-    if hasattr(gradtape, '_watch_accessed_variables'):
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent, gradtape._watch_accessed_variables)
+    if isinstance(optimizer, _LegacyOptimizer):
+        return _DistributedOptimizer(optimizer, name, use_locking, device_dense,
+                                     device_sparse, compression, sparse_as_dense)
+    elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
+        import horovod.tensorflow.keras as hvd_k
+        return hvd_k.DistributedOptimizer(optimizer, name, device_dense, device_sparse,
+                                          compression, sparse_as_dense)
     else:
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent)
+        raise ValueError('Provided optimizer doesn\'t inherit from either legacy '
+                         'TensorFlow or Keras optimizer: %s' % optimizer)
+
+
+if hasattr(tf, 'GradientTape'):
+    class _DistributedGradientTape(tf.GradientTape):
+        def __init__(self, tape, device_dense, device_sparse, compression, sparse_as_dense,
+                     persistent=False, watch_accessed_variables=True):
+            if hasattr(tape, '_watch_accessed_variables'):
+                super(self.__class__, self).__init__(persistent, watch_accessed_variables)
+            else:
+                super(self.__class__, self).__init__(persistent)
+
+            self._tape = tape
+            self._allreduce_grads = _make_allreduce_grads_fn(
+                'DistributedGradientTape', device_dense, device_sparse, compression,
+                sparse_as_dense)
+
+        def gradient(self, target, sources, output_gradients=None):
+            gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
+            if size() > 1:
+                return self._allreduce_grads(gradients)
+            else:
+                return gradients
+
+
+    def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
+                                compression=Compression.none, sparse_as_dense=False):
+        """An tape that wraps another tf.GradientTape, using an allreduce to
+        average gradient values before applying gradients to model weights.
+
+        Args:
+          gradtape:
+            GradientTape to use for computing gradients and applying updates.
+          device_dense:
+            Device to be used for dense tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLREDUCE.
path: horovod/tensorflow/__init__.py
text: @user1:
"build" -> "built"

diff_hunk: @@ -304,13 +273,69 @@ def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
         performance and memory utilization if the original sparse gradient
         has high density.  Defaults to false.
     """
-    cls = type(gradtape.__class__.__name__, (gradtape.__class__,),
-               dict(_DistributedGradientTape.__dict__))
-    if hasattr(gradtape, '_watch_accessed_variables'):
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent, gradtape._watch_accessed_variables)
+    if isinstance(optimizer, _LegacyOptimizer):
+        return _DistributedOptimizer(optimizer, name, use_locking, device_dense,
+                                     device_sparse, compression, sparse_as_dense)
+    elif isinstance(optimizer, tf.keras.optimizers.Optimizer):
+        import horovod.tensorflow.keras as hvd_k
+        return hvd_k.DistributedOptimizer(optimizer, name, device_dense, device_sparse,
+                                          compression, sparse_as_dense)
     else:
-        return cls(gradtape._tape, device_dense, device_sparse,
-                   compression, sparse_as_dense,
-                   gradtape._persistent)
+        raise ValueError('Provided optimizer doesn\'t inherit from either legacy '
+                         'TensorFlow or Keras optimizer: %s' % optimizer)
+
+
+if hasattr(tf, 'GradientTape'):
+    class _DistributedGradientTape(tf.GradientTape):
+        def __init__(self, tape, device_dense, device_sparse, compression, sparse_as_dense,
+                     persistent=False, watch_accessed_variables=True):
+            if hasattr(tape, '_watch_accessed_variables'):
+                super(self.__class__, self).__init__(persistent, watch_accessed_variables)
+            else:
+                super(self.__class__, self).__init__(persistent)
+
+            self._tape = tape
+            self._allreduce_grads = _make_allreduce_grads_fn(
+                'DistributedGradientTape', device_dense, device_sparse, compression,
+                sparse_as_dense)
+
+        def gradient(self, target, sources, output_gradients=None):
+            gradients = super(self.__class__, self).gradient(target, sources, output_gradients)
+            if size() > 1:
+                return self._allreduce_grads(gradients)
+            else:
+                return gradients
+
+
+    def DistributedGradientTape(gradtape, device_dense='', device_sparse='',
+                                compression=Compression.none, sparse_as_dense=False):
+        """An tape that wraps another tf.GradientTape, using an allreduce to
+        average gradient values before applying gradients to model weights.
+
+        Args:
+          gradtape:
+            GradientTape to use for computing gradients and applying updates.
+          device_dense:
+            Device to be used for dense tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLREDUCE.
+          device_sparse:
+            Device to be used for sparse tensors. Uses GPU by default
+            if Horovod was build with HOROVOD_GPU_ALLGATHER.
path: horovod/tensorflow/__init__.py
text: @user1:
"build" -> "built"

diff_hunk: @@ -0,0 +1,86 @@
+#!/usr/bin/env python
path: examples/tensorflow2_mnist.py
text: @user1:
Does this script have the executable bit set?

@author:
We don't have executable set on any of the examples.  I'll remove this `env python`.