Rename call sites from tff.Computation to `federated_language.Compu…

…tation`.

Match: tff\.Computation
Replace: federated_language.Computation
PiperOrigin-RevId: 711812790

docs/federated_learning.md

@@ -209,9 +209,9 @@ tutorials.
 
 ## Federated Computation Builders
 
-The `tff.learning` package provides several builders for `tff.Computation`s that
-perform learning-related tasks; we expect the set of such computations to expand
-in the future.
+The `tff.learning` package provides several builders for
+`federated_language.Computation`s that perform learning-related tasks; we expect
+the set of such computations to expand in the future.
 
 ### Architectural assumptions
 
@@ -250,8 +250,8 @@ The *federated computations* represented in this serialized form are expressed
 in a platform-independent internal language distinct from Python, but to use the
 Federated Learning API, you won't need to concern yourself with the details of
 this representation. The computations are represented in your Python code as
-objects of type `tff.Computation`, which for the most part you can treat as
-opaque Python `callable`s.
+objects of type `federated_language.Computation`, which for the most part you
+can treat as opaque Python `callable`s.
 
 In the tutorials, you will invoke those federated computations as if they were
 regular Python functions, to be executed locally. However, TFF is designed to

docs/tutorials/private_heavy_hitters.ipynb

@@ -98,6 +98,7 @@
         "\n",
         "# tensorflow_federated_nightly also bring in tf_nightly, which\n",
         "# can causes a duplicate tensorboard install, leading to errors.\n",
+        "!pip install --quite --upgrade federated_language\n",
         "!pip install --quiet tensorflow-text-nightly\n",
         "!pip install --quiet --upgrade tensorflow-federated"
       ]
@@ -295,7 +296,7 @@
       },
       "outputs": [],
       "source": [
-        "def run_simulation(one_round_computation: tff.Computation, dataset):\n",
+        "def run_simulation(one_round_computation: federated_language.Computation, dataset):\n",
         "  output = one_round_computation(dataset)\n",
         "  heavy_hitters = output.heavy_hitters\n",
         "  heavy_hitters_counts = output.heavy_hitters_counts\n",

docs/tutorials/working_with_client_data.ipynb

@@ -281,9 +281,9 @@
         "id": "Ek7W3ZZHMr1k"
       },
       "source": [
-        "## Interfacing with a `tff.Computation`\n",
+        "## Interfacing with a `federated_language.Computation`\n",
         "\n",
-        "Now that we can perform some basic manipulations with `ClientData` objects, we are ready to feed data to a `tff.Computation`. We define a [`tff.templates.IterativeProcess`](https://www.tensorflow.org/federated/api_docs/python/tff/templates/IterativeProcess) which implements [Federated Averaging](https://arxiv.org/abs/1602.05629), and explore different methods of passing it data."
+        "Now that we can perform some basic manipulations with `ClientData` objects, we are ready to feed data to a `federated_language.Computation`. We define a [`tff.templates.IterativeProcess`](https://www.tensorflow.org/federated/api_docs/python/tff/templates/IterativeProcess) which implements [Federated Averaging](https://arxiv.org/abs/1602.05629), and explore different methods of passing it data."
       ]
     },
     {
@@ -412,7 +412,7 @@
         "\n",
         "Conceptually, the reason for preferring this is clear: in TFF's local runtime, the clients only \"accidentally\" have access to the global Python environment due to the fact that the entire federated orchestration is happening on a single machine. It is worthwhile noting at this point that similar thinking gives rise to TFF's cross-platform, always-serializable, functional philosophy.\n",
         "\n",
-        "TFF makes such a change simple via `ClientData's` attribute `dataset_computation`, a `tff.Computation` which takes a `client_id` and returns the associated `tf.data.Dataset`.\n",
+        "TFF makes such a change simple via `ClientData's` attribute `dataset_computation`, a `federated_language.Computation` which takes a `client_id` and returns the associated `tf.data.Dataset`.\n",
         "\n",
         "Note that `preprocess` simply works with `dataset_computation`; the `dataset_computation` attribute of the preprocessed `ClientData` incorporates the entire preprocessing pipeline we just defined:"
       ]
@@ -552,7 +552,7 @@
         "id": "pHG0NXbWQuk7"
       },
       "source": [
-        "In fact, this is effectively what the helper we used is doing under the hood (plus performing appropriate type checking and manipulation). We could even have expressed the same logic slightly differently, by serializing `preprocess_and_shuffle` into a `tff.Computation`, and decomposing the `federated_map` into one step which constructs un-preprocessed datasets and another which runs `preprocess_and_shuffle` at each client.\n",
+        "In fact, this is effectively what the helper we used is doing under the hood (plus performing appropriate type checking and manipulation). We could even have expressed the same logic slightly differently, by serializing `preprocess_and_shuffle` into a `federated_language.Computation`, and decomposing the `federated_map` into one step which constructs un-preprocessed datasets and another which runs `preprocess_and_shuffle` at each client.\n",
         "\n",
         " We can verify that this more-manual path results in computations with the same type signature as TFF's helper (modulo parameter names):"
       ]

examples/personalization/emnist_p13n_main.py

@@ -190,8 +190,9 @@ def model_fn() -> tff.learning.models.VariableModel:
       num_epochs_per_eval=1,
   )
 
-  # Build the `tff.Computation` for evaluating the personalization strategies.
-  # Here `p13n_eval` is a `tff.Computation` with the following type signature:
+  # Build the `federated_language.Computation` for evaluating the
+  # personalization strategies. Here `p13n_eval` is a
+  # `federated_language.Computation` with the following type signature:
   # <model_weights@SERVER, datasets@CLIENTS> -> personalization_metrics@SERVER.
   p13n_eval = tff.learning.algorithms.build_personalization_eval_computation(
       model_fn=model_fn,

examples/program/BUILD

@@ -19,10 +19,7 @@ py_library(
 py_library(
     name = "program_logic",
     srcs = ["program_logic.py"],
-    deps = [
-        "//tensorflow_federated",
-        "@federated_language//federated_language",
-    ],
+    deps = ["@federated_language//federated_language"],
 )
 
 py_binary(

examples/program/program_logic.py

@@ -29,7 +29,6 @@
 from typing import NamedTuple, Optional
 
 import federated_language
-import tensorflow_federated as tff
 
 
 class UnexpectedTypeSignatureError(Exception):
@@ -38,10 +37,10 @@ class UnexpectedTypeSignatureError(Exception):
 
 def _check_expected_type_signatures(
     *,
-    initialize: tff.Computation,
-    train: tff.Computation,
+    initialize: federated_language.Computation,
+    train: federated_language.Computation,
     train_data_source: federated_language.program.FederatedDataSource,
-    evaluation: tff.Computation,
+    evaluation: federated_language.Computation,
     evaluation_data_source: federated_language.program.FederatedDataSource,
 ) -> None:
   """Checks the computations and data sources for the expected type signatures.
@@ -50,18 +49,18 @@ def _check_expected_type_signatures(
   example, if you are using a `tff.learning.templates.LearningProcess` as an
   input to the program logic, then these checks might not make sense because the
   the `tff.learning.templates.LearningProcess` has already validated that those
-  `tff.Computation` have the expected type signatures.
+  `federated_language.Computation` have the expected type signatures.
 
   See `train_federated_model` for more information on the expected type
   signatures of the computations and data sources.
 
   Args:
-    initialize: A `tff.Computation` to invoke before training.
-    train: A `tff.Computation` to invoke during training.
+    initialize: A `federated_language.Computation` to invoke before training.
+    train: A `federated_language.Computation` to invoke during training.
     train_data_source: A `federated_language.program.FederatedDataSource` which
       returns client data used during training.
-    evaluation: A `tff.Computation` to invoke to evaluate the model produced
-      after training.
+    evaluation: A `federated_language.Computation` to invoke to evaluate the
+      model produced after training.
     evaluation_data_source: A `federated_language.program.FederatedDataSource`
       which returns client data used during evaluation.
 
@@ -380,10 +379,10 @@ class _ProgramState(NamedTuple):
 
 async def train_federated_model(
     *,
-    initialize: tff.Computation,
-    train: tff.Computation,
+    initialize: federated_language.Computation,
+    train: federated_language.Computation,
     train_data_source: federated_language.program.FederatedDataSource,
-    evaluation: tff.Computation,
+    evaluation: federated_language.Computation,
     evaluation_data_source: federated_language.program.FederatedDataSource,
     total_rounds: int,
     num_clients: int,
@@ -449,12 +448,12 @@ async def train_federated_model(
   Finally, `state` is released to the `model_output_manager`.
 
   Args:
-    initialize: A `tff.Computation` to invoke before training.
-    train: A `tff.Computation` to invoke during training.
+    initialize: A `federated_language.Computation` to invoke before training.
+    train: A `federated_language.Computation` to invoke during training.
     train_data_source: A `federated_language.program.FederatedDataSource` which
       returns client data used during training.
-    evaluation: A `tff.Computation` to invoke to evaluate the model produced
-      after training.
+    evaluation: A `federated_language.Computation` to invoke to evaluate the
+      model produced after training.
     evaluation_data_source: A `federated_language.program.FederatedDataSource`
       which returns client data used during evaluation.
     total_rounds: The number of training rounds to run.

examples/program/program_logic_test.py

@@ -87,7 +87,7 @@ def _create_mock_initialize(
     side_effect: Optional[object] = None,
 ) -> mock.Mock:
   mock_initialize = mock.create_autospec(
-      tff.Computation, spec_set=True, side_effect=side_effect
+      federated_language.Computation, spec_set=True, side_effect=side_effect
   )
   type_signature = federated_language.FunctionType(None, state_type)
   type(mock_initialize).type_signature = mock.PropertyMock(
@@ -104,7 +104,7 @@ def _create_mock_train(
     side_effect: Optional[object] = None,
 ) -> mock.Mock:
   mock_train = mock.create_autospec(
-      tff.Computation, spec_set=True, side_effect=side_effect
+      federated_language.Computation, spec_set=True, side_effect=side_effect
   )
   type_signature = federated_language.FunctionType(
       [state_type, train_data_type],
@@ -124,7 +124,7 @@ def _create_mock_evaluation(
     side_effect: Optional[object] = None,
 ) -> mock.Mock:
   mock_evaluation = mock.create_autospec(
-      tff.Computation, spec_set=True, side_effect=side_effect
+      federated_language.Computation, spec_set=True, side_effect=side_effect
   )
   type_signature = federated_language.FunctionType(
       [state_type, evaluation_data_type], evaluation_metrics_type

tensorflow_federated/python/aggregators/encoded.py

@@ -152,7 +152,8 @@ def _encoded_init_fn(encoders):
     encoders: A collection of `GatherEncoder` objects.
 
   Returns:
-    A no-arg `tff.Computation` returning initial state for `EncodedSumFactory`.
+    A no-arg `federated_language.Computation` returning initial state for
+    `EncodedSumFactory`.
   """
   init_fn_tf = tensorflow_computation.tf_computation(
       lambda: tf.nest.map_structure(lambda e: e.initial_state(), encoders)
@@ -184,7 +185,8 @@ def _encoded_next_fn(server_state_type, value_type, encoders):
     encoders: A collection of `GatherEncoder` objects.
 
   Returns:
-    A `tff.Computation` for `EncodedSumFactory`, with the type signature of
+    A `federated_language.Computation` for `EncodedSumFactory`, with the type
+    signature of
     `(server_state_type, value_type@CLIENTS) ->
     MeasuredProcessOutput(server_state_type, value_type@SERVER, ()@SERVER)`
   """

tensorflow_federated/python/aggregators/rotation.py

@@ -374,7 +374,8 @@ def _build_next_global_seed_fn(stride):
     stride: An integer increment to the counter in the seed.
 
   Returns:
-    A `tff.Computation` that takes and returns `tf.int64` tensor with shape [2].
+    A `federated_language.Computation` that takes and returns `tf.int64` tensor
+    with shape [2].
   """
 
   @tensorflow_computation.tf_computation(SEED_TFF_TYPE)

tensorflow_federated/python/analytics/count_distinct.py

@@ -52,13 +52,14 @@ def _log2(u: tf.Tensor) -> tf.Tensor:
 def build_client_hyperloglog_computation() -> (
     federated_language.framework.Computation
 ):
-  """Builds a `tff.Computation` for computing client hyperloglog sketches.
+  """Builds a `federated_language.Computation` for computing client hyperloglog sketches.
 
   Specifically, the returned computation consumes a dataset of integer hashes
   and returns the HyperLogLog sketch of size HLL_SKETCH_SIZE.
 
   Returns:
-    A `tff.Computation` for computing client hyperloglog sketches.
+    A `federated_language.Computation` for computing client hyperloglog
+    sketches.
   """
 
   @tensorflow_computation.tf_computation(
@@ -91,7 +92,7 @@ def reduce_func(state, hash_value):
 def build_federated_secure_max_computation() -> (
     federated_language.framework.Computation
 ):
-  """Builds a `tff.Computation` for computing max in a secure fashion.
+  """Builds a `federated_language.Computation` for computing max in a secure fashion.
 
     Specifically, the returned computation consumes sketches at @CLIENTS and
     returns the element-wise max of the inpt sketches @SERVER.
@@ -105,7 +106,7 @@ def build_federated_secure_max_computation() -> (
     (upper_bound - lower_bound).
 
   Returns:
-    A `tff.Computation` for computing max of client vectors.
+    A `federated_language.Computation` for computing max of client vectors.
   """
 
   @federated_language.federated_computation(
@@ -162,7 +163,7 @@ def maxes_from_onehots(x):
 def create_federated_hyperloglog_computation(
     *, use_secagg: bool = False
 ) -> federated_language.framework.Computation:
-  """Creates a `tff.Computation` to estimate the number of distinct strings.
+  """Creates a `federated_language.Computation` to estimate the number of distinct strings.
 
   The returned computation consumes data @CLIENTS and produces an estimate of
   the number of unique words across all clients @SERVER.
@@ -172,7 +173,7 @@ def create_federated_hyperloglog_computation(
       computing the hyperloglog sketch.
 
   Returns:
-    A `tff.Computation` for running the HyperLogLog algorithm.
+    A `federated_language.Computation` for running the HyperLogLog algorithm.
   """
 
   @tensorflow_computation.tf_computation

tensorflow_federated/python/analytics/heavy_hitters/iblt/iblt_tff.py

@@ -73,7 +73,7 @@ def build_iblt_computation(
     batch_size: int = 1,
     repetitions: int = 3,
 ) -> federated_language.framework.Computation:
-  """Builds the `tff.Computation` for heavy-hitters discovery with IBLT.
+  """Builds the `federated_language.Computation` for heavy-hitters discovery with IBLT.
 
   Args:
     string_max_bytes: The maximum length in bytes of a string in the IBLT.
@@ -137,7 +137,8 @@ def build_iblt_computation(
       override the default. Defaults to `3`. Must be at least `3`.
 
   Returns:
-    A `tff.Computation` that performs federated heavy hitter discovery.
+    A `federated_language.Computation` that performs federated heavy hitter
+    discovery.
 
   Raises:
     ValueError: if parameters don't meet expectations.

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

@@ -151,7 +151,7 @@ def round_comp(server_state, client_data):
   client population. In some systems, all accumulation may happen at the
   server, but for scalability reasons, it is often desirable to structure
   aggregation in multiple tiers. Its type signature is `A`, or when
-  represented as a `tff.Computation` in Python, `( -> A)`.
+  represented as a `federated_language.Computation` in Python, `( -> A)`.
 
 * `accumulate` is the TensorFlow computation that updates the state of an
   update accumulator (initialized with `zero` above) with a single client's

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

@@ -60,7 +60,7 @@
 def get_computation_for_broadcast_form(
     bf: forms.BroadcastForm,
 ) -> federated_language.framework.Computation:
-  """Creates `tff.Computation` from a broadcast form."""
+  """Creates `federated_language.Computation` from a broadcast form."""
   py_typecheck.check_type(bf, forms.BroadcastForm)
   server_data_type = bf.compute_server_context.type_signature.parameter
   client_data_type = bf.client_processing.type_signature.parameter[1]
@@ -155,13 +155,13 @@ def get_state_initialization_computation(
 def get_computation_for_map_reduce_form(
     mrf: forms.MapReduceForm,
 ) -> federated_language.framework.Computation:
-  """Creates `tff.Computation` from a MapReduce form.
+  """Creates `federated_language.Computation` from a MapReduce form.
 
   Args:
     mrf: An instance of `tff.backends.mapreduce.MapReduceForm`.
 
   Returns:
-    An instance of `tff.Computation` that corresponds to `mrf`.
+    An instance of `federated_language.Computation` that corresponds to `mrf`.
 
   Raises:
     TypeError: If the arguments are of the wrong types.
@@ -219,13 +219,13 @@ def computation(arg):
 def get_computation_for_distribute_aggregate_form(
     daf: forms.DistributeAggregateForm,
 ) -> federated_language.framework.Computation:
-  """Creates `tff.Computation` from a DistributeAggregate form.
+  """Creates `federated_language.Computation` from a DistributeAggregate form.
 
   Args:
     daf: An instance of `tff.backends.mapreduce.DistributeAggregateForm`.
 
   Returns:
-    An instance of `tff.Computation` that corresponds to `daf`.
+    An instance of `federated_language.Computation` that corresponds to `daf`.
 
   Raises:
     TypeError: If the arguments are of the wrong types.
@@ -1012,7 +1012,7 @@ def get_broadcast_form_for_computation(
 
   Returns:
     An instance of `tff.backends.mapreduce.BroadcastForm` equivalent to the
-    provided `tff.Computation`.
+    provided `federated_language.Computation`.
   """
   py_typecheck.check_type(
       comp, federated_language.framework.ConcreteComputation