Skip to content

Commit d7dc221

Browse files
michaelreneercopybara-github
michaelreneer
authored and
copybara-github
committed
Move some TensorFlow compiler functions to the TensorFlow backend.
PiperOrigin-RevId: 647449336
1 parent b33947c commit d7dc221

18 files changed

+72
-72
lines changed

‎tensorflow_federated/python/core/backends/mapreduce/BUILD‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -63,11 +63,11 @@ py_test(
6363
":form_utils",
6464
":mapreduce_test_utils",
6565
"//tensorflow_federated/proto/v0:computation_py_pb2",
66+
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
6667
"//tensorflow_federated/python/core/impl/compiler:building_block_factory",
6768
"//tensorflow_federated/python/core/impl/compiler:building_block_test_utils",
6869
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
6970
"//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
70-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
7171
"//tensorflow_federated/python/core/impl/compiler:transformation_utils",
7272
"//tensorflow_federated/python/core/impl/compiler:tree_analysis",
7373
"//tensorflow_federated/python/core/impl/computation:computation_impl",
@@ -194,10 +194,10 @@ py_library(
194194
deps = [
195195
"//tensorflow_federated/python/common_libs:py_typecheck",
196196
"//tensorflow_federated/python/common_libs:structure",
197+
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
197198
"//tensorflow_federated/python/core/impl/compiler:building_block_factory",
198199
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
199200
"//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
200-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
201201
"//tensorflow_federated/python/core/impl/context_stack:context_base",
202202
"//tensorflow_federated/python/core/impl/context_stack:context_stack_impl",
203203
"//tensorflow_federated/python/core/impl/context_stack:symbol_binding_context",

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,11 @@
2020
from tensorflow_federated.python.core.backends.mapreduce import compiler
2121
from tensorflow_federated.python.core.backends.mapreduce import form_utils
2222
from tensorflow_federated.python.core.backends.mapreduce import mapreduce_test_utils
23+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
2324
from tensorflow_federated.python.core.impl.compiler import building_block_factory
2425
from tensorflow_federated.python.core.impl.compiler import building_block_test_utils
2526
from tensorflow_federated.python.core.impl.compiler import building_blocks
2627
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
27-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
2828
from tensorflow_federated.python.core.impl.compiler import transformation_utils
2929
from tensorflow_federated.python.core.impl.compiler import tree_analysis
3030
from tensorflow_federated.python.core.impl.computation import computation_impl

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,10 @@
1717

1818
from tensorflow_federated.python.common_libs import py_typecheck
1919
from tensorflow_federated.python.common_libs import structure
20+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
2021
from tensorflow_federated.python.core.impl.compiler import building_block_factory
2122
from tensorflow_federated.python.core.impl.compiler import building_blocks
2223
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
23-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
2424
from tensorflow_federated.python.core.impl.context_stack import context_base
2525
from tensorflow_federated.python.core.impl.context_stack import context_stack_impl
2626
from tensorflow_federated.python.core.impl.context_stack import symbol_binding_context

‎tensorflow_federated/python/core/backends/native/BUILD‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -130,10 +130,10 @@ py_library(
130130
deps = [
131131
"//tensorflow_federated/python/core/backends/mapreduce:compiler",
132132
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_building_block_factory",
133+
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
133134
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_tree_transformations",
134135
"//tensorflow_federated/python/core/impl/compiler:building_block_factory",
135136
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
136-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
137137
"//tensorflow_federated/python/core/impl/compiler:transformations",
138138
"//tensorflow_federated/python/core/impl/compiler:tree_analysis",
139139
"//tensorflow_federated/python/core/impl/compiler:tree_transformations",

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,10 +15,10 @@
1515

1616
from tensorflow_federated.python.core.backends.mapreduce import compiler
1717
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_building_block_factory
18+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
1819
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_tree_transformations
1920
from tensorflow_federated.python.core.impl.compiler import building_block_factory
2021
from tensorflow_federated.python.core.impl.compiler import building_blocks
21-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
2222
from tensorflow_federated.python.core.impl.compiler import transformations
2323
from tensorflow_federated.python.core.impl.compiler import tree_analysis
2424
from tensorflow_federated.python.core.impl.compiler import tree_transformations

‎tensorflow_federated/python/core/backends/test/BUILD‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,11 +36,11 @@ py_library(
3636
"//tensorflow_federated/python/common_libs:structure",
3737
"//tensorflow_federated/python/core/backends/mapreduce:intrinsics",
3838
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_building_block_factory",
39+
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_computation_factory",
3940
"//tensorflow_federated/python/core/environments/tensorflow_backend:tensorflow_tree_transformations",
4041
"//tensorflow_federated/python/core/impl/compiler:building_block_factory",
4142
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
4243
"//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
43-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
4444
"//tensorflow_federated/python/core/impl/computation:computation_impl",
4545
"//tensorflow_federated/python/core/impl/types:computation_types",
4646
"//tensorflow_federated/python/core/impl/types:placements",

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,11 +23,11 @@
2323
from tensorflow_federated.python.common_libs import structure
2424
from tensorflow_federated.python.core.backends.mapreduce import intrinsics as mapreduce_intrinsics
2525
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_building_block_factory
26+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
2627
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_tree_transformations
2728
from tensorflow_federated.python.core.impl.compiler import building_block_factory
2829
from tensorflow_federated.python.core.impl.compiler import building_blocks
2930
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
30-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
3131
from tensorflow_federated.python.core.impl.computation import computation_impl
3232
from tensorflow_federated.python.core.impl.types import computation_types
3333
from tensorflow_federated.python.core.impl.types import placements

‎tensorflow_federated/python/core/environments/tensorflow_backend/BUILD‎

Lines changed: 55 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
load("@rules_python//python:defs.bzl", "py_library", "py_test")
2+
load("//tools:build_defs.bzl", "py_cpu_gpu_test")
23

34
package(
45
default_applicable_licenses = ["//:package_license"],
@@ -45,10 +46,10 @@ py_test(
4546
srcs = ["compiled_computation_transformations_test.py"],
4647
deps = [
4748
":compiled_computation_transformations",
49+
":tensorflow_computation_factory",
50+
":tensorflow_computation_test_utils",
4851
":tensorflow_computation_transformations",
4952
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
50-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
51-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_test_utils",
5253
"//tensorflow_federated/python/core/impl/types:computation_types",
5354
],
5455
)
@@ -83,12 +84,12 @@ py_library(
8384
name = "tensorflow_building_block_factory",
8485
srcs = ["tensorflow_building_block_factory.py"],
8586
deps = [
87+
":tensorflow_computation_factory",
8688
"//tensorflow_federated/python/common_libs:py_typecheck",
8789
"//tensorflow_federated/python/common_libs:structure",
8890
"//tensorflow_federated/python/core/impl/compiler:building_block_factory",
8991
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
9092
"//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
91-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
9293
"//tensorflow_federated/python/core/impl/types:array_shape",
9394
"//tensorflow_federated/python/core/impl/types:computation_types",
9495
"//tensorflow_federated/python/core/impl/types:placements",
@@ -102,14 +103,61 @@ py_test(
102103
srcs = ["tensorflow_building_block_factory_test.py"],
103104
deps = [
104105
":tensorflow_building_block_factory",
106+
":tensorflow_computation_test_utils",
105107
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
106108
"//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
107-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_test_utils",
108109
"//tensorflow_federated/python/core/impl/types:computation_types",
109110
"//tensorflow_federated/python/core/impl/types:placements",
110111
],
111112
)
112113

114+
py_library(
115+
name = "tensorflow_computation_factory",
116+
srcs = ["tensorflow_computation_factory.py"],
117+
deps = [
118+
"//tensorflow_federated/proto/v0:computation_py_pb2",
119+
"//tensorflow_federated/python/common_libs:py_typecheck",
120+
"//tensorflow_federated/python/common_libs:structure",
121+
"//tensorflow_federated/python/core/impl/compiler:local_computation_factory_base",
122+
"//tensorflow_federated/python/core/impl/types:array_shape",
123+
"//tensorflow_federated/python/core/impl/types:computation_types",
124+
"//tensorflow_federated/python/core/impl/types:type_analysis",
125+
"//tensorflow_federated/python/core/impl/types:type_conversions",
126+
"//tensorflow_federated/python/core/impl/types:type_serialization",
127+
"//tensorflow_federated/python/core/impl/types:type_transformations",
128+
"//tensorflow_federated/python/core/impl/utils:tensorflow_utils",
129+
"//tensorflow_federated/python/tensorflow_libs:serialization_utils",
130+
],
131+
)
132+
133+
py_cpu_gpu_test(
134+
name = "tensorflow_computation_factory_test",
135+
srcs = ["tensorflow_computation_factory_test.py"],
136+
deps = [
137+
":tensorflow_computation_factory",
138+
":tensorflow_computation_test_utils",
139+
"//tensorflow_federated/proto/v0:computation_py_pb2",
140+
"//tensorflow_federated/python/common_libs:structure",
141+
"//tensorflow_federated/python/core/impl/types:computation_types",
142+
"//tensorflow_federated/python/core/impl/types:placements",
143+
"//tensorflow_federated/python/core/impl/types:type_factory",
144+
"//tensorflow_federated/python/core/impl/types:type_serialization",
145+
],
146+
)
147+
148+
py_library(
149+
name = "tensorflow_computation_test_utils",
150+
testonly = True,
151+
srcs = ["tensorflow_computation_test_utils.py"],
152+
deps = [
153+
"//tensorflow_federated/proto/v0:computation_py_pb2",
154+
"//tensorflow_federated/python/common_libs:structure",
155+
"//tensorflow_federated/python/core/impl/types:computation_types",
156+
"//tensorflow_federated/python/core/impl/types:type_serialization",
157+
"//tensorflow_federated/python/core/impl/utils:tensorflow_utils",
158+
],
159+
)
160+
113161
py_library(
114162
name = "tensorflow_computation_transformations",
115163
srcs = ["tensorflow_computation_transformations.py"],
@@ -125,10 +173,10 @@ py_test(
125173
size = "small",
126174
srcs = ["tensorflow_computation_transformations_test.py"],
127175
deps = [
176+
":tensorflow_computation_factory",
128177
":tensorflow_computation_transformations",
129178
"//tensorflow_federated/proto/v0:computation_py_pb2",
130179
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
131-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
132180
"//tensorflow_federated/python/core/impl/types:computation_types",
133181
"//tensorflow_federated/python/core/impl/types:type_serialization",
134182
"//tensorflow_federated/python/core/impl/utils:tensorflow_utils",
@@ -152,9 +200,9 @@ py_test(
152200
"nokokoro", # b/193543632: C++ execution is not fully supported in OSS.
153201
],
154202
deps = [
203+
":tensorflow_computation_factory",
155204
":tensorflow_executor_bindings",
156205
"//tensorflow_federated/proto/v0:executor_py_pb2",
157-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
158206
"//tensorflow_federated/python/core/impl/executors:executor_bindings",
159207
"//tensorflow_federated/python/core/impl/executors:value_serialization",
160208
"//tensorflow_federated/python/core/impl/types:computation_types",
@@ -168,11 +216,11 @@ py_library(
168216
srcs = ["tensorflow_tree_transformations.py"],
169217
deps = [
170218
":tensorflow_building_block_factory",
219+
":tensorflow_computation_factory",
171220
"//tensorflow_federated/python/common_libs:py_typecheck",
172221
"//tensorflow_federated/python/core/impl/compiler:building_block_factory",
173222
"//tensorflow_federated/python/core/impl/compiler:building_blocks",
174223
"//tensorflow_federated/python/core/impl/compiler:intrinsic_defs",
175-
"//tensorflow_federated/python/core/impl/compiler:tensorflow_computation_factory",
176224
"//tensorflow_federated/python/core/impl/compiler:transformation_utils",
177225
"//tensorflow_federated/python/core/impl/types:computation_types",
178226
"//tensorflow_federated/python/core/impl/types:type_analysis",

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,10 +17,10 @@
1717
import tensorflow as tf
1818

1919
from tensorflow_federated.python.core.environments.tensorflow_backend import compiled_computation_transformations
20+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
21+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_test_utils
2022
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_transformations
2123
from tensorflow_federated.python.core.impl.compiler import building_blocks
22-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
23-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_test_utils
2424
from tensorflow_federated.python.core.impl.types import computation_types
2525

2626

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@
1919

2020
from tensorflow_federated.python.common_libs import py_typecheck
2121
from tensorflow_federated.python.common_libs import structure
22+
from tensorflow_federated.python.core.environments.tensorflow_backend import tensorflow_computation_factory
2223
from tensorflow_federated.python.core.impl.compiler import building_block_factory
2324
from tensorflow_federated.python.core.impl.compiler import building_blocks
2425
from tensorflow_federated.python.core.impl.compiler import intrinsic_defs
25-
from tensorflow_federated.python.core.impl.compiler import tensorflow_computation_factory
2626
from tensorflow_federated.python.core.impl.types import array_shape
2727
from tensorflow_federated.python.core.impl.types import computation_types
2828
from tensorflow_federated.python.core.impl.types import placements

0 commit comments

Comments
 (0)