Remove the XLA dependency from the executors package.

PiperOrigin-RevId: 584650864
google-parfait · Nov 22, 2023 · 4fa6173 · 4fa6173
1 parent 21d9f3f
commit 4fa6173
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 10 deletions.
diff --git a/tensorflow_federated/python/core/impl/executors/BUILD b/tensorflow_federated/python/core/impl/executors/BUILD
@@ -177,7 +177,6 @@ py_test(
     deps = [
         ":executor_bindings",
         ":value_serialization",
-        ":xla_executor_bindings",
         "//tensorflow_federated/proto/v0:executor_py_pb2",
         "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
         "//tensorflow_federated/python/core/impl/types:computation_types",

diff --git a/tensorflow_federated/python/core/impl/executors/executor_bindings_test.py b/tensorflow_federated/python/core/impl/executors/executor_bindings_test.py
@@ -25,7 +25,6 @@
 from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.executors import executor_bindings
 from tensorflow_federated.python.core.impl.executors import value_serialization
-from tensorflow_federated.python.core.impl.executors import xla_executor_bindings
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import placements
 from tensorflow_federated.python.core.impl.types import type_conversions
@@ -705,17 +704,15 @@ def test_roundtrip(self, input_value, dtype):
 class SequenceExecutorBindingsTest(absltest.TestCase):
 
   def test_create(self):
+    executor = executor_bindings.create_tensorflow_executor()
     try:
-      executor_bindings.create_sequence_executor(
-          xla_executor_bindings.create_xla_executor()
-      )
-    except Exception as e:  # pylint: disable=broad-except
-      self.fail(f'Exception: {e}')
+      executor_bindings.create_sequence_executor(executor)
+    except Exception:  # pylint: disable=broad-except
+      self.fail('Raised `Exception` unexpectedly.')
 
   def test_materialize_on_unkown_fails(self):
-    executor = executor_bindings.create_sequence_executor(
-        xla_executor_bindings.create_xla_executor()
-    )
+    executor = executor_bindings.create_tensorflow_executor()
+    executor_bindings.create_sequence_executor(executor)
     with self.assertRaisesRegex(Exception, 'NOT_FOUND'):
       executor.materialize(0)