xla::CustomCallSchedule requires TF2.6.

Signed-off-by: Trent Lo <trentl@nvidia.com>
horovod · Aug 11, 2021 · 72df0ae · 72df0ae
1 parent 2337d81
commit 72df0ae
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/horovod/tensorflow/xla_mpi_ops.cc b/horovod/tensorflow/xla_mpi_ops.cc
@@ -19,7 +19,7 @@
 #include <thread>
 #include <unordered_map>
 
-#if TENSORFLOW_VERSION >= 2005000000
+#if TENSORFLOW_VERSION >= 2006000000
 
 #include "tensorflow/compiler/tf2xla/xla_op_kernel.h"
 #include "tensorflow/compiler/tf2xla/xla_op_registry.h"
@@ -570,6 +570,6 @@ XLA_REGISTER_CUSTOM_CALL_TARGET(CallbackHVDAllreduceDone, "CUDA");
 } // namespace tensorflow
 } // namespace horovod
 
-#endif // TENSORFLOW_VERSION >= 2005000000
+#endif // TENSORFLOW_VERSION >= 2006000000
 #endif // HAVE_CUDA
 #endif // HAVE_GPU
diff --git a/test/parallel/test_xla.py b/test/parallel/test_xla.py
@@ -60,14 +60,14 @@
 ccl_supported_types = set([tf.uint8, tf.int8, tf.uint16, tf.int16,
                            tf.int32, tf.int64, tf.float32])
 
-_IS_TF25 = LooseVersion(tf.__version__) >= LooseVersion('2.5.0')
+_IS_TF26 = LooseVersion(tf.__version__) >= LooseVersion('2.6.0')
 
 # Set environment variable to enable adding/removing process sets after
 # initializing Horovod.
 os.environ["HOROVOD_DYNAMIC_PROCESS_SETS"] = "1"
 
 
-@pytest.mark.skipif(not _IS_TF25, reason='TF2.5+ is required')
+@pytest.mark.skipif(not _IS_TF26, reason='TF2.6+ is required')
 class XLATests(tf.test.TestCase):
     """
     Tests for ops in horovod.tensorflow.