Move some TensorFlow compiler functions to the TensorFlow backend.

PiperOrigin-RevId: 647449336

tensorflow_federated/python/core/backends/mapreduce/BUILD

@@ -63,11 +63,11 @@ py_test(
  ":form_utils",
  ":mapreduce_test_utils",
  "//tensorflow_federated/proto/v0:computation_py_pb2",
+ "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_block_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_block_test_utils",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
  "//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/compiler:transformation_utils",
  "//tensorflow_federated/python/core/impl/compiler:tree_analysis",
  "//tensorflow_federated/python/core/impl/computation:computation_impl",
@@ -194,10 +194,10 @@ py_library(
  deps = [
  "//tensorflow_federated/python/common_libs:py_typecheck",
  "//tensorflow_federated/python/common_libs:structure",
+ "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_block_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
  "//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/context_stack:context_base",
  "//tensorflow_federated/python/core/impl/context_stack:context_stack_impl",
  "//tensorflow_federated/python/core/impl/context_stack:symbol_binding_context",

tensorflow_federated/python/core/backends/mapreduce/compiler_test.py

@@ -20,11 +20,11 @@
 from tensorflow_federated.python.core.backends.mapreduce import compiler
 from tensorflow_federated.python.core.backends.mapreduce import form_utils
 from tensorflow_federated.python.core.backends.mapreduce import mapreduce_test_utils
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import building_block_factory
 from tensorflow_federated.python.core.impl.compiler import building_block_test_utils
 from tensorflow_federated.python.core.impl.compiler import building_blocks
 from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import transformation_utils
 from tensorflow_federated.python.core.impl.compiler import tree_analysis
 from tensorflow_federated.python.core.impl.computation import computation_impl

tensorflow_federated/python/core/backends/mapreduce/intrinsics.py

@@ -17,10 +17,10 @@
 
 from tensorflow_federated.python.common_libs import py_typecheck
 from tensorflow_federated.python.common_libs import structure
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import building_block_factory
 from tensorflow_federated.python.core.impl.compiler import building_blocks
 from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.context_stack import context_base
 from tensorflow_federated.python.core.impl.context_stack import context_stack_impl
 from tensorflow_federated.python.core.impl.context_stack import symbol_binding_context

tensorflow_federated/python/core/backends/native/BUILD

@@ -130,10 +130,10 @@ py_library(
  deps = [
  "//tensorflow_federated/python/core/backends/mapreduce:compiler",
  "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_building_block_factory",
+ "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_tree_transformations",
  "//tensorflow_federated/python/core/impl/compiler:building_block_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/compiler:transformations",
  "//tensorflow_federated/python/core/impl/compiler:tree_analysis",
  "//tensorflow_federated/python/core/impl/compiler:tree_transformations",

tensorflow_federated/python/core/backends/native/mergeable_comp_compiler.py

@@ -15,10 +15,10 @@
 
 from tensorflow_federated.python.core.backends.mapreduce import compiler
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_building_block_factory
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_tree_transformations
 from tensorflow_federated.python.core.impl.compiler import building_block_factory
 from tensorflow_federated.python.core.impl.compiler import building_blocks
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import transformations
 from tensorflow_federated.python.core.impl.compiler import tree_analysis
 from tensorflow_federated.python.core.impl.compiler import tree_transformations

tensorflow_federated/python/core/backends/test/BUILD

@@ -36,11 +36,11 @@ py_library(
  "//tensorflow_federated/python/common_libs:structure",
  "//tensorflow_federated/python/core/backends/mapreduce:intrinsics",
  "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_building_block_factory",
+ "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_tree_transformations",
  "//tensorflow_federated/python/core/impl/compiler:building_block_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
  "//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/computation:computation_impl",
  "//tensorflow_federated/python/core/impl/types:computation_types",
  "//tensorflow_federated/python/core/impl/types:placements",

tensorflow_federated/python/core/backends/test/compiler.py

@@ -23,11 +23,11 @@
 from tensorflow_federated.python.common_libs import structure
 from tensorflow_federated.python.core.backends.mapreduce import intrinsics as mapreduce_intrinsics
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_building_block_factory
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_tree_transformations
 from tensorflow_federated.python.core.impl.compiler import building_block_factory
 from tensorflow_federated.python.core.impl.compiler import building_blocks
 from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.computation import computation_impl
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import placements

tensorflow_federated/python/core/environments/tensorflow_backend/BUILD

@@ -1,4 +1,5 @@
 load("@rules_python//python:defs.bzl", "py_library", "py_test")
+load("//tools:build_defs.bzl", "py_cpu_gpu_test")
 
 package(
  default_applicable_licenses = ["//:package_license"],
@@ -45,10 +46,10 @@ py_test(
  srcs = ["compiled_computation_transformations_test.py"],
  deps = [
  ":compiled_computation_transformations",
+ ":tensorflow_computation_factory",
+ ":tensorflow_computation_test_utils",
  ":tensorflow_computation_transformations",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_test_utils",
  "//tensorflow_federated/python/core/impl/types:computation_types",
  ],
 )
@@ -83,12 +84,12 @@ py_library(
  name = "tensorflow_building_block_factory",
  srcs = ["tensorflow_building_block_factory.py"],
  deps = [
+ ":tensorflow_computation_factory",
  "//tensorflow_federated/python/common_libs:py_typecheck",
  "//tensorflow_federated/python/common_libs:structure",
  "//tensorflow_federated/python/core/impl/compiler:building_block_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
  "//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/types:array_shape",
  "//tensorflow_federated/python/core/impl/types:computation_types",
  "//tensorflow_federated/python/core/impl/types:placements",
@@ -102,14 +103,61 @@ py_test(
  srcs = ["tensorflow_building_block_factory_test.py"],
  deps = [
  ":tensorflow_building_block_factory",
+ ":tensorflow_computation_test_utils",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
  "//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_test_utils",
  "//tensorflow_federated/python/core/impl/types:computation_types",
  "//tensorflow_federated/python/core/impl/types:placements",
  ],
 )
 
+py_library(
+ name = "tensorflow_computation_factory",
+ srcs = ["tensorflow_computation_factory.py"],
+ deps = [
+ "//tensorflow_federated/proto/v0:computation_py_pb2",
+ "//tensorflow_federated/python/common_libs:py_typecheck",
+ "//tensorflow_federated/python/common_libs:structure",
+ "//tensorflow_federated/python/core/impl/compiler:local_computation_factory_base",
+ "//tensorflow_federated/python/core/impl/types:array_shape",
+ "//tensorflow_federated/python/core/impl/types:computation_types",
+ "//tensorflow_federated/python/core/impl/types:type_analysis",
+ "//tensorflow_federated/python/core/impl/types:type_conversions",
+ "//tensorflow_federated/python/core/impl/types:type_serialization",
+ "//tensorflow_federated/python/core/impl/types:type_transformations",
+ "//tensorflow_federated/python/core/impl/utils:tensorflow_utils",
+ "//tensorflow_federated/python/tensorflow_libs:serialization_utils",
+ ],
+)
+
+py_cpu_gpu_test(
+ name = "tensorflow_computation_factory_test",
+ srcs = ["tensorflow_computation_factory_test.py"],
+ deps = [
+ ":tensorflow_computation_factory",
+ ":tensorflow_computation_test_utils",
+ "//tensorflow_federated/proto/v0:computation_py_pb2",
+ "//tensorflow_federated/python/common_libs:structure",
+ "//tensorflow_federated/python/core/impl/types:computation_types",
+ "//tensorflow_federated/python/core/impl/types:placements",
+ "//tensorflow_federated/python/core/impl/types:type_factory",
+ "//tensorflow_federated/python/core/impl/types:type_serialization",
+ ],
+)
+
+py_library(
+ name = "tensorflow_computation_test_utils",
+ testonly = True,
+ srcs = ["tensorflow_computation_test_utils.py"],
+ deps = [
+ "//tensorflow_federated/proto/v0:computation_py_pb2",
+ "//tensorflow_federated/python/common_libs:structure",
+ "//tensorflow_federated/python/core/impl/types:computation_types",
+ "//tensorflow_federated/python/core/impl/types:type_serialization",
+ "//tensorflow_federated/python/core/impl/utils:tensorflow_utils",
+ ],
+)
+
 py_library(
  name = "tensorflow_computation_transformations",
  srcs = ["tensorflow_computation_transformations.py"],
@@ -125,10 +173,10 @@ py_test(
  size = "small",
  srcs = ["tensorflow_computation_transformations_test.py"],
  deps = [
+ ":tensorflow_computation_factory",
  ":tensorflow_computation_transformations",
  "//tensorflow_federated/proto/v0:computation_py_pb2",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/types:computation_types",
  "//tensorflow_federated/python/core/impl/types:type_serialization",
  "//tensorflow_federated/python/core/impl/utils:tensorflow_utils",
@@ -152,9 +200,9 @@ py_test(
  "nokokoro", # b/193543632: C++ execution is not fully supported in OSS.
  ],
  deps = [
+ ":tensorflow_computation_factory",
  ":tensorflow_executor_bindings",
  "//tensorflow_federated/proto/v0:executor_py_pb2",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/executors:executor_bindings",
  "//tensorflow_federated/python/core/impl/executors:value_serialization",
  "//tensorflow_federated/python/core/impl/types:computation_types",
@@ -168,11 +216,11 @@ py_library(
  srcs = ["tensorflow_tree_transformations.py"],
  deps = [
  ":tensorflow_building_block_factory",
+ ":tensorflow_computation_factory",
  "//tensorflow_federated/python/common_libs:py_typecheck",
  "//tensorflow_federated/python/core/impl/compiler:building_block_factory",
  "//tensorflow_federated/python/core/impl/compiler:building_blocks",
  "//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
- "//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
  "//tensorflow_federated/python/core/impl/compiler:transformation_utils",
  "//tensorflow_federated/python/core/impl/types:computation_types",
  "//tensorflow_federated/python/core/impl/types:type_analysis",

tensorflow_federated/python/core/environments/tensorflow_backend/compiled_computation_transformations_test.py

@@ -17,10 +17,10 @@
 import tensorflow as tf
 
 from tensorflow_federated.python.core.environments.tensorflow_backend import compiled_computation_transformations
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_test_utils
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_transformations
 from tensorflow_federated.python.core.impl.compiler import building_blocks
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_test_utils
 from tensorflow_federated.python.core.impl.types import computation_types
 
 

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_building_block_factory.py

@@ -19,10 +19,10 @@
 
 from tensorflow_federated.python.common_libs import py_typecheck
 from tensorflow_federated.python.common_libs import structure
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import building_block_factory
 from tensorflow_federated.python.core.impl.compiler import building_blocks
 from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.types import array_shape
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import placements

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_building_block_factory_test.py

@@ -17,9 +17,9 @@
 import numpy as np
 
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_building_block_factory
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_test_utils
 from tensorflow_federated.python.core.impl.compiler import building_blocks
 from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_test_utils
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import placements
 

tensorflow_federated/python/core/impl/compiler/tensorflow_computation_factory_test.py → tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_computation_factory_test.py

@@ -21,8 +21,8 @@
 
 from tensorflow_federated.proto.v0 import computation_pb2 as pb
 from tensorflow_federated.python.common_libs import structure
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_test_utils
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_test_utils
 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_factory

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_computation_transformations_test.py

@@ -20,9 +20,9 @@
 import tensorflow as tf
 
 from tensorflow_federated.proto.v0 import computation_pb2
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_transformations
 from tensorflow_federated.python.core.impl.compiler import building_blocks
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import type_serialization
 from tensorflow_federated.python.core.impl.utils import tensorflow_utils

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_executor_bindings_test.py

@@ -20,8 +20,8 @@
 import tensorflow as tf
 
 from tensorflow_federated.proto.v0 import executor_pb2
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_executor_bindings
-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.types import computation_types

tensorflow_federated/python/core/environments/tensorflow_backend/tensorflow_tree_transformations.py

@@ -21,10 +21,10 @@
 
 from tensorflow_federated.python.common_libs import py_typecheck
 from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_building_block_factory
+from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import building_block_factory
 from tensorflow_federated.python.core.impl.compiler import building_blocks
 from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
-from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
 from tensorflow_federated.python.core.impl.compiler import transformation_utils
 from tensorflow_federated.python.core.impl.types import computation_types
 from tensorflow_federated.python.core.impl.types import type_analysis