From ce93275c020da29438ede1569c489dee3d7145d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Chlo=C3=A9=20Kiddon?= <loeki@google.com>
Date: Mon, 6 Jan 2025 09:24:56 -0800
Subject: [PATCH] Migrate tff.federated_computation to
 federated_language.federated_computation. ...in .py code locations.

This API has moved to the new federated_language project.

PiperOrigin-RevId: 712553861
---
 examples/program/computations.py              |  6 +--
 examples/simple_fedavg/simple_fedavg_tff.py   |  6 +--
 .../stateful_clients/stateful_fedavg_tff.py   |  4 +-
 .../python/aggregators/primitives_test.py     |  6 ++-
 .../core/backends/mapreduce/__init__.py       |  4 +-
 .../python/core/backends/mapreduce/forms.py   |  3 +-
 .../mergeable_comp_execution_context.py       |  8 ++--
 .../python/learning/metrics/aggregator.py     |  9 ++--
 tensorflow_federated/python/tests/BUILD       |  5 ++-
 .../python/tests/ast_generation_test.py       |  5 ++-
 ...sync_execution_context_integration_test.py |  8 ++--
 .../python/tests/backend_accelerators_test.py |  4 +-
 .../python/tests/backend_test.py              | 45 ++++++++++---------
 ...comp_execution_context_integration_test.py | 37 ++++++++-------
 .../remote_runtime_stream_structs_test.py     |  2 +-
 .../sync_local_cpp_execution_context_test.py  |  4 +-
 .../tests/temperature_sensor_example.py       |  2 +-
 tools/python_package/test_python_package.sh   |  1 -
 18 files changed, 88 insertions(+), 71 deletions(-)

diff --git a/examples/program/computations.py b/examples/program/computations.py
index 4e348b576b..6f26aab954 100644
--- a/examples/program/computations.py
+++ b/examples/program/computations.py
@@ -26,7 +26,7 @@
 METRICS_TOTAL_SUM = 'total_sum'
 
 
-@tff.federated_computation()
+@federated_language.federated_computation()
 def initialize():
   """Returns the initial state."""
   return tff.federated_value(0, federated_language.SERVER)
@@ -44,7 +44,7 @@ def _sum_integers(x: int, y: int) -> int:
   return x + y
 
 
-@tff.federated_computation(
+@federated_language.federated_computation(
     federated_language.FederatedType(np.int32, federated_language.SERVER),
     federated_language.FederatedType(
         federated_language.SequenceType(np.int32), federated_language.CLIENTS
@@ -78,7 +78,7 @@ def train(server_state: int, client_data: tf.data.Dataset):
   return updated_state, metrics
 
 
-@tff.federated_computation(
+@federated_language.federated_computation(
     federated_language.FederatedType(np.int32, federated_language.SERVER),
     federated_language.FederatedType(
         federated_language.SequenceType(np.int32), federated_language.CLIENTS
diff --git a/examples/simple_fedavg/simple_fedavg_tff.py b/examples/simple_fedavg/simple_fedavg_tff.py
index 783454ad55..3b51c16aa9 100644
--- a/examples/simple_fedavg/simple_fedavg_tff.py
+++ b/examples/simple_fedavg/simple_fedavg_tff.py
@@ -96,7 +96,7 @@ def initialize_client_data(server_message):
       model = model_fn()
       return init_client_ouput(model, server_message)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         server_message_type, federated_language.SequenceType(batch_type)
     )
     def client_update_weights_fn(server_message, batches):
@@ -174,7 +174,7 @@ def server_message_fn(server_state):
       tf_dataset_type, federated_language.CLIENTS
   )
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       federated_server_state_type, federated_dataset_type
   )
   def run_one_round(server_state, federated_dataset):
@@ -221,7 +221,7 @@ def run_one_round(server_state, federated_dataset):
     )
     return server_state, aggregated_outputs
 
-  @tff.federated_computation
+  @federated_language.federated_computation
   def server_init_tff():
     """Orchestration logic for server model initialization."""
     return tff.federated_eval(server_init_tf, federated_language.SERVER)
diff --git a/examples/stateful_clients/stateful_fedavg_tff.py b/examples/stateful_clients/stateful_fedavg_tff.py
index dd800d7fdb..d0966c1b98 100644
--- a/examples/stateful_clients/stateful_fedavg_tff.py
+++ b/examples/stateful_clients/stateful_fedavg_tff.py
@@ -137,7 +137,7 @@ def client_update_fn(tf_dataset, client_state, server_message):
       client_state_type, federated_language.CLIENTS
   )
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       federated_server_state_type,
       federated_dataset_type,
       federated_client_state_type,
@@ -178,7 +178,7 @@ def run_one_round(server_state, federated_dataset, client_states):
 
     return server_state, round_loss_metric, client_outputs.client_state
 
-  @tff.federated_computation
+  @federated_language.federated_computation
   def server_init_tff():
     """Orchestration logic for server model initialization."""
     return tff.federated_value(server_init_tf(), federated_language.SERVER)
diff --git a/tensorflow_federated/python/aggregators/primitives_test.py b/tensorflow_federated/python/aggregators/primitives_test.py
index d387a28878..170d82d295 100644
--- a/tensorflow_federated/python/aggregators/primitives_test.py
+++ b/tensorflow_federated/python/aggregators/primitives_test.py
@@ -1254,7 +1254,8 @@ def _build_test_sum_fn_py_bounds(value_type, lower_bound, upper_bound):
     upper_bound: A Python numeric constant or a numpy array.
 
   Returns:
-    A `tff.federated_computation` with type signature `(value_type@CLIENTS ->
+    A `federated_language.federated_computation` with type signature
+    `(value_type@CLIENTS ->
     value_type@SERVER)`.
   """
 
@@ -1286,7 +1287,8 @@ def _build_test_sum_fn_tff_bounds(
     upper_bound_type: A `federated_language.Type` of upper_bound to be used.
 
   Returns:
-    A `tff.federated_computation` with type signature `((value_type@CLIENTS,
+    A `federated_language.federated_computation` with type signature
+    `((value_type@CLIENTS,
     lower_bound_type@SERVER, upper_bound_type@SERVER) -> value_type@SERVER)`.
   """
 
diff --git a/tensorflow_federated/python/core/backends/mapreduce/__init__.py b/tensorflow_federated/python/core/backends/mapreduce/__init__.py
index a9085f3557..fb4061184a 100644
--- a/tensorflow_federated/python/core/backends/mapreduce/__init__.py
+++ b/tensorflow_federated/python/core/backends/mapreduce/__init__.py
@@ -80,7 +80,7 @@
 `MapReduceForm` maps to a single federated round.
 
 ```python
-@tff.federated_computation
+@federated_language.federated_computation
 def round_comp(server_state, client_data):
 
   # The server prepares an input to be broadcast to all clients that controls
@@ -208,7 +208,7 @@ def round_comp(server_state, client_data):
 single federated round.
 
 ```python
-@tff.federated_computation
+@federated_language.federated_computation
 def round_comp(server_state, client_data):
   # The server prepares an input to be broadcast to all clients and generates
   # a temporary state that may be used by later parts of the computation.
diff --git a/tensorflow_federated/python/core/backends/mapreduce/forms.py b/tensorflow_federated/python/core/backends/mapreduce/forms.py
index cdb6a8b64f..d35ad94788 100644
--- a/tensorflow_federated/python/core/backends/mapreduce/forms.py
+++ b/tensorflow_federated/python/core/backends/mapreduce/forms.py
@@ -36,6 +36,7 @@ def _check_tensorflow_computation(label, comp):
 
 
 def _check_lambda_computation(label, comp):
+  """Validates a Lambda computation."""
   py_typecheck.check_type(
       comp, federated_language.framework.ConcreteComputation, label
   )
@@ -131,7 +132,7 @@ class BroadcastForm:
   ```
   server_data_type = self.compute_server_context.type_signature.parameter
   client_data_type = self.client_processing.type_signature.parameter[1]
-  @tff.federated_computation(server_data_type, client_data_type)
+  @federated_language.federated_computation(server_data_type, client_data_type)
   def _(server_data, client_data):
     # Select out the bit of server context to send to the clients.
     context_at_server = tff.federated_map(
diff --git a/tensorflow_federated/python/core/impl/execution_contexts/mergeable_comp_execution_context.py b/tensorflow_federated/python/core/impl/execution_contexts/mergeable_comp_execution_context.py
index 9fbcd20f36..fae0edf677 100644
--- a/tensorflow_federated/python/core/impl/execution_contexts/mergeable_comp_execution_context.py
+++ b/tensorflow_federated/python/core/impl/execution_contexts/mergeable_comp_execution_context.py
@@ -70,7 +70,7 @@ class MergeableCompForm:
   without changing the results of the computation, a computation of the form:
 
   ```
-  @tff.federated_computation(...)
+  @federated_language.federated_computation(...)
   def single_aggregation(arg):
     result_at_clients = work(arg)
     agg_result = tff.federated_aggregate(
@@ -80,19 +80,19 @@ def single_aggregation(arg):
 
   can be represented as the `MergeableCompForm` triplet:
   ```
-  @tff.federated_computation(federated_language.AbstractType('T'))
+  @federated_language.federated_computation(federated_language.AbstractType('T'))
   def up_to_merge(arg):
     result_at_clients = work(arg)
     agg_result = tff.federated_aggregate(
         result_at_clients, accumulate_zero, accumulate, merge, identity_report)
     return agg_result
 
-  @tff.federated_computation([up_to_merge.type_signature.result.member,
+  @federated_language.federated_computation([up_to_merge.type_signature.result.member,
                               up_to_merge.type_signature.result.member])
   def merge(arg):
     return merge(arg[0], arg[1])
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       federated_language.AbstractType('T'),
       federated_language.FederatedType(merge.type_signature.result,
       federated_language.SERVER),
diff --git a/tensorflow_federated/python/learning/metrics/aggregator.py b/tensorflow_federated/python/learning/metrics/aggregator.py
index 008fb6b13e..84a24592ba 100644
--- a/tensorflow_federated/python/learning/metrics/aggregator.py
+++ b/tensorflow_federated/python/learning/metrics/aggregator.py
@@ -50,7 +50,8 @@ def sum_then_finalize(
   the polymorphic method is invoked on.
 
   Note: invoking this computation outside of a federated context (a method
-  decorated with `tff.federated_computation`) will require first wrapping it in
+  decorated with `federated_language.federated_computation`) will require first
+  wrapping it in
   a concrete, non-polymorphic `federated_language.Computation` with appropriate
   federated
   types.
@@ -148,7 +149,8 @@ def secure_sum_then_finalize(
   the polymorphic method is invoked on.
 
   Note: invoking this computation outside of a federated context (a method
-  decorated with `tff.federated_computation`) will require first wrapping it in
+  decorated with `federated_language.federated_computation`) will require first
+  wrapping it in
   a concrete, non-polymorphic `federated_language.Computation` with appropriate
   federated
   types.
@@ -319,7 +321,8 @@ def finalize_then_sample(
   the polymorphic method is invoked on.
 
   Note: invoking this computation outside of a federated context (a method
-  decoratedc with `tff.federated_computation`) will require first wrapping it in
+  decoratedc with `federated_language.federated_computation`) will require first
+  wrapping it in
   a concrete, non-polymorphic `federated_language.Computation` with appropriate
   federated
   types.
diff --git a/tensorflow_federated/python/tests/BUILD b/tensorflow_federated/python/tests/BUILD
index 72cd9fa1ab..23c9eec7c3 100644
--- a/tensorflow_federated/python/tests/BUILD
+++ b/tensorflow_federated/python/tests/BUILD
@@ -17,7 +17,10 @@ py_test(
     name = "ast_generation_test",
     size = "small",
     srcs = ["ast_generation_test.py"],
-    deps = ["//tensorflow_federated"],
+    deps = [
+        "//tensorflow_federated",
+        "@federated_language//federated_language",
+    ],
 )
 
 py_test(
diff --git a/tensorflow_federated/python/tests/ast_generation_test.py b/tensorflow_federated/python/tests/ast_generation_test.py
index 4dfe43f0e7..0987fa501e 100644
--- a/tensorflow_federated/python/tests/ast_generation_test.py
+++ b/tensorflow_federated/python/tests/ast_generation_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 from absl.testing import absltest
+import federated_language
 import tensorflow as tf
 import tensorflow_federated as tff
 
@@ -36,7 +37,7 @@ def test_flattens_to_tf_computation(self):
     def five():
       return 5
 
-    @tff.federated_computation
+    @federated_language.federated_computation
     def federated_five():
       return five()
 
@@ -56,7 +57,7 @@ def test_only_one_random_only_generates_a_single_call_to_random(self):
     def rand():
       return tf.random.normal([])
 
-    @tff.federated_computation
+    @federated_language.federated_computation
     def same_rand_tuple():
       single_random_number = rand()
       return (single_random_number, single_random_number)
diff --git a/tensorflow_federated/python/tests/async_execution_context_integration_test.py b/tensorflow_federated/python/tests/async_execution_context_integration_test.py
index 951f5e728c..2043b885ca 100644
--- a/tensorflow_federated/python/tests/async_execution_context_integration_test.py
+++ b/tensorflow_federated/python/tests/async_execution_context_integration_test.py
@@ -28,7 +28,7 @@ async def test_install_and_execute_in_context(self):
     factory = tff.framework.local_cpp_executor_factory()
     context = tff.framework.AsyncExecutionContext(factory)
 
-    @tff.federated_computation(np.int32)
+    @federated_language.federated_computation(np.int32)
     def identity(x):
       return x
 
@@ -42,7 +42,7 @@ async def test_install_and_execute_computations_with_different_cardinalities(
     factory = tff.framework.local_cpp_executor_factory()
     context = tff.framework.AsyncExecutionContext(factory)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.int32, federated_language.CLIENTS)
     )
     def repackage_arg(x):
@@ -60,7 +60,7 @@ async def test_runs_cardinality_free(self):
         factory, cardinality_inference_fn=(lambda x, y: {})
     )
 
-    @tff.federated_computation(np.int32)
+    @federated_language.federated_computation(np.int32)
     def identity(x):
       return x
 
@@ -85,7 +85,7 @@ def _cardinality_fn(x, y):
         np.int32, federated_language.CLIENTS
     )
 
-    @tff.federated_computation(arg_type)
+    @federated_language.federated_computation(arg_type)
     def identity(x):
       return x
 
diff --git a/tensorflow_federated/python/tests/backend_accelerators_test.py b/tensorflow_federated/python/tests/backend_accelerators_test.py
index eaf2e94a7d..3237dd7a13 100644
--- a/tensorflow_federated/python/tests/backend_accelerators_test.py
+++ b/tensorflow_federated/python/tests/backend_accelerators_test.py
@@ -34,7 +34,7 @@ def dataset_reduce_fn_wrapper(ds):
     initial_val = tf.Variable(np.int64(1.0))
     return dataset_reduce_fn(ds, initial_val)
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       tff.at_clients(federated_language.SequenceType(np.int64))
   )
   def parallel_client_run(client_datasets):
@@ -59,7 +59,7 @@ def dataset_reduce_fn_wrapper(ds):
     initial_val = tf.Variable(np.int64(1.0))
     return dataset_reduce_fn(ds, initial_val)
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       tff.at_clients(federated_language.SequenceType(np.int64))
   )
   def parallel_client_run(client_datasets):
diff --git a/tensorflow_federated/python/tests/backend_test.py b/tensorflow_federated/python/tests/backend_test.py
index d65b155fce..1f2d51b614 100644
--- a/tensorflow_federated/python/tests/backend_test.py
+++ b/tensorflow_federated/python/tests/backend_test.py
@@ -40,7 +40,7 @@ def accumulate(a, b):
     def report(a):
       return a
 
-    @tff.federated_computation()
+    @federated_language.federated_computation()
     def empty_agg():
       val_at_clients = tff.federated_value([], federated_language.CLIENTS)
       return tff.federated_aggregate(
@@ -53,7 +53,7 @@ def empty_agg():
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_executes_empty_sum(self):
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.int32, federated_language.CLIENTS)
     )
     def fed_sum(x):
@@ -68,7 +68,7 @@ def test_empty_mean_returns_nan(self):
     # TODO: b/200970992 - Standardize handling of this case. We currently have a
     # ZeroDivisionError, a RuntimeError, and a context that returns nan.
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.float32, federated_language.CLIENTS)
     )
     def fed_mean(x):
@@ -107,7 +107,7 @@ def concat_datasets(ds1, ds2):
     def identity(ds):
       return ds
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             federated_language.SequenceType(element_type),
             federated_language.CLIENTS,
@@ -163,7 +163,8 @@ class FederatedComputationTest(parameterized.TestCase):
 
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_constant(self):
-    @tff.federated_computation
+
+    @federated_language.federated_computation
     def foo():
       return 10
 
@@ -172,7 +173,8 @@ def foo():
 
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_empty_tuple(self):
-    @tff.federated_computation
+
+    @federated_language.federated_computation
     def foo():
       return ()
 
@@ -181,7 +183,8 @@ def foo():
 
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_federated_value(self):
-    @tff.federated_computation
+
+    @federated_language.federated_computation
     def foo(x):
       return tff.federated_value(x, federated_language.SERVER)
 
@@ -191,7 +194,7 @@ def foo(x):
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_federated_zip(self):
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         [federated_language.FederatedType(np.int32, federated_language.CLIENTS)]
         * 2
     )
@@ -208,7 +211,7 @@ def test_federated_zip_with_twenty_elements(self):
     num_element = 20
     num_clients = 2
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         [federated_language.FederatedType(np.int32, federated_language.CLIENTS)]
         * num_element
     )
@@ -226,7 +229,7 @@ def test_repeated_invocations_of_map(self):
     def add_one(x):
       return x + 1
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.int32, federated_language.CLIENTS)
     )
     def map_add_one(federated_arg):
@@ -245,7 +248,7 @@ def test_polymorphism(self):
     def add_one(x):
       return x + 1
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.int32, federated_language.CLIENTS)
     )
     def map_add_one(federated_arg):
@@ -263,7 +266,7 @@ def map_add_one(federated_arg):
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_runs_unplaced_lambda(self):
 
-    @tff.federated_computation(np.int32, np.int32)
+    @federated_language.federated_computation(np.int32, np.int32)
     def bar(x, y):
       del y  # Unused
       return x
@@ -274,12 +277,12 @@ def bar(x, y):
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_runs_server_placed_lambda(self):
 
-    @tff.federated_computation(np.int32, np.int32)
+    @federated_language.federated_computation(np.int32, np.int32)
     def foo(x, y):
       del y  # Unused
       return x
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             collections.OrderedDict(x=np.int32, y=np.int32),
             federated_language.SERVER,
@@ -294,12 +297,12 @@ def bar(server_tuple):
   @tff.test.with_contexts(*test_contexts.get_all_contexts())
   def test_runs_clients_placed_lambda(self):
 
-    @tff.federated_computation(np.int32, np.int32)
+    @federated_language.federated_computation(np.int32, np.int32)
     def foo(x, y):
       del y  # Unused
       return x
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             collections.OrderedDict(x=np.int32, y=np.int32),
             federated_language.CLIENTS,
@@ -321,7 +324,7 @@ def foo(x):
       # reshape.
       return tf.reshape(x, [])
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             tensor_type, federated_language.CLIENTS
         )
@@ -329,7 +332,7 @@ def foo(x):
     def map_foo_at_clients(x):
       return tff.federated_map(foo, x)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(tensor_type, federated_language.SERVER)
     )
     def map_foo_at_server(x):
@@ -361,7 +364,7 @@ def test_runs_federated_select(self):
     def gather(selectee, key):
       return tf.gather(selectee, key)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             selectee_type, federated_language.SERVER
         ),
@@ -541,7 +544,7 @@ def test_computation_called_once_is_invoked_once(self):
     def get_random():
       return tf.random.normal([])
 
-    @tff.federated_computation
+    @federated_language.federated_computation
     def get_one_random_twice():
       value = get_random()
       return value, value
@@ -556,7 +559,7 @@ def test_computation_called_twice_is_invoked_twice(self):
     def get_random():
       return tf.random.normal([])
 
-    @tff.federated_computation
+    @federated_language.federated_computation
     def get_two_random():
       return get_random(), get_random()
 
diff --git a/tensorflow_federated/python/tests/mergeable_comp_execution_context_integration_test.py b/tensorflow_federated/python/tests/mergeable_comp_execution_context_integration_test.py
index f65dc6a805..b0f75aa406 100644
--- a/tensorflow_federated/python/tests/mergeable_comp_execution_context_integration_test.py
+++ b/tensorflow_federated/python/tests/mergeable_comp_execution_context_integration_test.py
@@ -25,7 +25,8 @@ def build_sum_client_arg_computation(
     server_arg_type: federated_language.FederatedType,
     clients_arg_type: federated_language.FederatedType,
 ) -> federated_language.Computation:
-  @tff.federated_computation(server_arg_type, clients_arg_type)
+
+  @federated_language.federated_computation(server_arg_type, clients_arg_type)
   def up_to_merge(server_arg, client_arg):
     del server_arg  # Unused
     return tff.federated_sum(client_arg)
@@ -34,7 +35,8 @@ def up_to_merge(server_arg, client_arg):
 
 
 def build_noarg_count_clients_computation() -> federated_language.Computation:
-  @tff.federated_computation()
+
+  @federated_language.federated_computation()
   def up_to_merge():
     return tff.federated_sum(tff.federated_value(1, federated_language.CLIENTS))
 
@@ -44,7 +46,8 @@ def up_to_merge():
 def build_whimsy_merge_computation(
     arg_type: federated_language.Type,
 ) -> federated_language.Computation:
-  @tff.federated_computation(arg_type, arg_type)
+
+  @federated_language.federated_computation(arg_type, arg_type)
   def merge(arg0, arg1):
     del arg1  # Unused
     return arg0
@@ -69,7 +72,7 @@ def build_whimsy_after_merge_computation(
 ) -> federated_language.Computation:
   if original_arg_type is not None:
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         original_arg_type,
         federated_language.FederatedType(
             merge_result_type, federated_language.SERVER
@@ -81,7 +84,7 @@ def after_merge(original_arg, merge_result):
 
   else:
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             merge_result_type, federated_language.SERVER
         )
@@ -97,7 +100,7 @@ def build_return_merge_result_computation(
     merge_result_type: federated_language.Type,
 ) -> federated_language.Computation:
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       original_arg_type,
       federated_language.FederatedType(
           merge_result_type, federated_language.SERVER
@@ -114,7 +117,7 @@ def build_return_merge_result_with_no_first_arg_computation(
     merge_result_type: federated_language.Type,
 ) -> federated_language.Computation:
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       federated_language.FederatedType(
           merge_result_type, federated_language.SERVER
       )
@@ -135,7 +138,7 @@ def build_sum_merge_with_first_arg_computation(
   def add(x, y):
     return x + y
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       original_arg_type,
       federated_language.FederatedType(
           merge_result_type, federated_language.SERVER
@@ -157,7 +160,7 @@ def test_raises_mismatched_up_to_merge_and_merge(self):
 
     bad_merge = build_whimsy_merge_computation(np.float32)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         up_to_merge.type_signature.parameter,
         federated_language.FederatedType(
             bad_merge.type_signature.result, federated_language.SERVER
@@ -177,12 +180,12 @@ def test_raises_merge_computation_not_assignable_result(self):
         federated_language.FederatedType(np.int32, federated_language.CLIENTS),
     )
 
-    @tff.federated_computation(np.int32, np.int32)
+    @federated_language.federated_computation(np.int32, np.int32)
     def bad_merge(x, y):
       del x, y  # Unused
       return 1.0  # of type float.
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         up_to_merge.type_signature.parameter,
         federated_language.FederatedType(
             bad_merge.type_signature.result, federated_language.SERVER
@@ -204,7 +207,7 @@ def test_raises_no_top_level_argument_in_after_agg(self):
 
     merge = build_whimsy_merge_computation(np.int32)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(
             merge.type_signature.result, federated_language.SERVER
         )
@@ -219,7 +222,7 @@ def bad_after_merge(x):
 
   def test_raises_up_to_merge_returns_non_server_placed_result(self):
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.int32, federated_language.SERVER)
     )
     def bad_up_to_merge(x):
@@ -245,7 +248,7 @@ def test_raises_with_aggregation_in_after_agg(self):
 
     merge = build_whimsy_merge_computation(np.int32)
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         up_to_merge.type_signature.parameter,
         federated_language.FederatedType(
             merge.type_signature.result, federated_language.SERVER
@@ -548,7 +551,8 @@ def test_counts_clients_with_noarg_computation(self, num_subrounds):
     self.assertEqual(result, expected_result)
 
   def test_invoke_raises_computation_no_compiler(self):
-    @tff.federated_computation()
+
+    @federated_language.federated_computation()
     def return_one():
       return 1
 
@@ -564,7 +568,8 @@ def return_one():
       context.invoke(return_one)
 
   def test_invoke_raises_computation_not_compiled_to_mergeable_comp_form(self):
-    @tff.federated_computation()
+
+    @federated_language.federated_computation()
     def return_one():
       return 1
 
diff --git a/tensorflow_federated/python/tests/remote_runtime_stream_structs_test.py b/tensorflow_federated/python/tests/remote_runtime_stream_structs_test.py
index 0d28aaa14a..19d8d97938 100644
--- a/tensorflow_federated/python/tests/remote_runtime_stream_structs_test.py
+++ b/tensorflow_federated/python/tests/remote_runtime_stream_structs_test.py
@@ -38,7 +38,7 @@ def _make_federated(
 ) -> federated_language.Computation:
   """Construct a federate computation that maps comptuation to CLIENTS."""
 
-  @tff.federated_computation(
+  @federated_language.federated_computation(
       federated_language.FederatedType(
           computation.type_signature.parameter, federated_language.CLIENTS
       ),
diff --git a/tensorflow_federated/python/tests/sync_local_cpp_execution_context_test.py b/tensorflow_federated/python/tests/sync_local_cpp_execution_context_test.py
index 23c6ddce32..9316ac9dac 100644
--- a/tensorflow_federated/python/tests/sync_local_cpp_execution_context_test.py
+++ b/tensorflow_federated/python/tests/sync_local_cpp_execution_context_test.py
@@ -99,7 +99,7 @@ def comp():
   )
   def test_changing_cardinalities_across_calls(self):
 
-    @tff.federated_computation(
+    @federated_language.federated_computation(
         federated_language.FederatedType(np.int32, federated_language.CLIENTS)
     )
     def comp(x):
@@ -119,7 +119,7 @@ def comp(x):
   )
   def test_conflicting_cardinalities_within_call(self):
 
-    @tff.federated_computation([
+    @federated_language.federated_computation([
         federated_language.FederatedType(np.int32, federated_language.CLIENTS),
         federated_language.FederatedType(np.int32, federated_language.CLIENTS),
     ])
diff --git a/tensorflow_federated/python/tests/temperature_sensor_example.py b/tensorflow_federated/python/tests/temperature_sensor_example.py
index 80af6e5484..40cf31e7cc 100644
--- a/tensorflow_federated/python/tests/temperature_sensor_example.py
+++ b/tensorflow_federated/python/tests/temperature_sensor_example.py
@@ -33,7 +33,7 @@ def count_total(ds):
   return ds.reduce(np.float32(0.0), lambda n, _: n + 1.0)
 
 
-@tff.federated_computation(
+@federated_language.federated_computation(
     federated_language.FederatedType(
         federated_language.SequenceType(np.float32), federated_language.CLIENTS
     ),
diff --git a/tools/python_package/test_python_package.sh b/tools/python_package/test_python_package.sh
index 137fa2deae..6c92fc8769 100755
--- a/tools/python_package/test_python_package.sh
+++ b/tools/python_package/test_python_package.sh
@@ -73,7 +73,6 @@ main() {
   # Test the Python package.
   pip install --upgrade "${package}"
   pip freeze
-  python -c "import tensorflow_federated as tff; print(tff.federated_computation(lambda: 'Hello World')())"
   python -c "import tensorflow_federated as tff; print(tff.__version__)"
 }