docs/tutorials/custom_aggregators.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "a930wM_fqUNH"
      },
      "source": [
        "##### Copyright 2021 The TensorFlow Federated Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "VxVUPYkahDa6"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j-y8GPCrZ--r"
      },
      "source": [
        "# Implementing Custom Aggregations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4jM4S9YFXamd"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/federated/tutorials/custom_aggregators\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/federated/blob/v0.76.0/docs/tutorials/custom_aggregators.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/federated/blob/v0.76.0/docs/tutorials/custom_aggregators.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/federated/docs/tutorials/custom_aggregators.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xcrxeoeXaWJC"
      },
      "source": [
        "In this tutorial, we explain design principles behind the `tff.aggregators` module and best practices for implementing custom aggregation of values from clients to server.\n",
        "\n",
        "**Prerequisites.** This tutorial assumes you are already familiar with basic concepts of [Federated Core](https://www.tensorflow.org/federated/federated_core) such as placements (`tff.SERVER`, `tff.CLIENTS`), how TFF represents computations (`tff.tf_computation`, `tff.federated_computation`) and their type signatures."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "MJqkeeDphf38"
      },
      "outputs": [],
      "source": [
        "#@test {\"skip\": true}\n",
        "!pip install --quiet --upgrade tensorflow-federated"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_WbNgHj7aRzm"
      },
      "source": [
        "## Design summary"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "essB1Vu0aXp1"
      },
      "source": [
        "In TFF, \"aggregation\" refers to the movement of a set of values on `tff.CLIENTS` to produce an aggregate value of the same type on `tff.SERVER`. That is, each individual client value need not be available. For example in federated learning, client model updates are averaged to get an aggregate model update to apply to the global model on the server.\n",
        "\n",
        "In addition to operators accomplishing this goal such as `tff.federated_sum`, TFF provides `tff.templates.AggregationProcess` (a [stateful process](https://www.tensorflow.org/federated/federated_learning#modeling_state)) which formalizes the type signature for aggregation computation so it can generalize to more complex forms than a simple sum.\n",
        "\n",
        "The main components of the `tff.aggregators` module are *factories* for creation of the `AggregationProcess`, which are designed to be generally useful and replacable building blocks of TFF in two aspects:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XmoO7KLSEAI3"
      },
      "source": [
        "1. *Parameterized computations.* Aggregation is an independent building block that can be plugged into other TFF modules designed to work with `tff.aggregators` to parameterize their necessary aggregation.\n",
        "\n",
        "Example:\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SNs6eeLOEIS_"
      },
      "source": [
        "```\n",
        "learning_process = tff.learning.algorithms.build_weighted_fed_avg(\n",
        "    ...,\n",
        "    model_aggregator=tff.aggregators.MeanFactory())\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nIXwjousD9XE"
      },
      "source": [
        "2. *Aggregation composition.* An aggregation building block can be composed with other aggregation building blocks to create more complex composite aggregations.\n",
        "\n",
        "Example:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "h94jEnY1bmhZ"
      },
      "source": [
        "```\n",
        "secure_mean = tff.aggregators.MeanFactory(\n",
        "    value_sum_factory=tff.aggregators.SecureSumFactory(...))\n",
        "```"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hH87D2OJEaAA"
      },
      "source": [
        "The rest of this tutorial explains how these two goals are achieved."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zUu7slQMfUX1"
      },
      "source": [
        "### Aggregation process"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wLn76zvHaYUh"
      },
      "source": [
        "We first summarize the `tff.templates.AggregationProcess`, and follow with the factory pattern for its creation.\n",
        "\n",
        "The `tff.templates.AggregationProcess` is an `tff.templates.MeasuredProcess` with type signatures specified for aggregation. In particular, the `initialize` and `next` functions have the following type signatures:\n",
        "\n",
        "* `( -\u003e state_type@SERVER)`\n",
        "* `(\u003cstate_type@SERVER, {value_type}@CLIENTS, *\u003e -\u003e \u003cstate_type@SERVER, value_type@SERVER, measurements_type@SERVER\u003e)`\n",
        "\n",
        "The state (of type `state_type`) must be placed at server. The `next` function takes as input argument the state and a value to be aggregated (of type `value_type`) placed at clients. The `*` means optional other input arguments, for instance weights in a weighted mean. It returns an updated state object, the aggregated value of the same type placed at server, and some measurements.\n",
        "\n",
        "Note that both the state to be passed between executions of the `next` function, and the reported measurements intended to report any information depending on a specific execution of the `next` function, may be empty. Nevertheless, they have to be explicitly specified for other parts of TFF to have a clear contract to follow.\n",
        "\n",
        "Other TFF modules, for instance the model updates in `tff.learning`, are expected to use the `tff.templates.AggregationProcess` to parameterize how values are aggregated. However, what exactly are the values aggregated and what their type signatures are, depends on other details of the model being trained and the learning algorithm used to do it.\n",
        "\n",
        "To make aggregation independent of the other aspects of computations, we use the factory pattern -- we create the appropriate `tff.templates.AggregationProcess` once the relevant type signatures of objects to be aggregated are available, by invoking the `create` method of the factory. Direct handling of the aggregation process is thus needed only for library authors, who are responsible for this creation."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dWoZQqpOgTv4"
      },
      "source": [
        "#### Aggregation process factories"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x6fh5D7jgXuX"
      },
      "source": [
        "There are two abstract base factory classes for unweighted and weighted aggregation. Their `create` method takes type signatures of value to be aggregated and returns a `tff.templates.AggregationProcess` for aggregation of such values.\n",
        "\n",
        "The process created by `tff.aggregators.UnweightedAggregationFactory` takes two input arguments: (1) state at server and (2) value of specified type `value_type`.\n",
        "\n",
        "An example implementation is `tff.aggregators.SumFactory`.\n",
        "\n",
        "The process created by `tff.aggregators.WeightedAggregationFactory` takes three input arguments: (1) state at server, (2) value of specified type `value_type` and (3) weight of type `weight_type`, as specified by the factory's user when invoking its `create` method.\n",
        "\n",
        "An example implementation is `tff.aggregators.MeanFactory` which computes a weighted mean.\n",
        "\n",
        "The factory pattern is how we achieve the first goal stated above; that aggregation is an independent building block. For example, when changing which model variables are trainable, a complex aggregation does not necessarily need to change; the factory representing it will be invoked with a different type signature when used by a method such as `tff.learning.algorithms.build_weighted_fed_avg`."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I8k46_iLgfSu"
      },
      "source": [
        "### Compositions\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZZPXuhSkgonh"
      },
      "source": [
        "Recall that a general aggregation process can encapsulate (a) some preprocessing of the values at clients, (b) movement of values from client to server, and (c) some postprocessing of the aggregated value at the server. The second goal stated above, aggregation composition, is realized inside the `tff.aggregators` module by structuring the implementation of the aggregation factories such that part (b) can be delegated to another aggregation factory.\n",
        "\n",
        "Rather than implementing all necessary logic within a single factory class, the implementations are by default focused on a single aspect relevant for aggregation. When needed, this pattern then enables us to replace the building blocks one at a time.\n",
        "\n",
        "An example is the weighted `tff.aggregators.MeanFactory`. Its implementation multiplies provided values and weights at clients, then sums both weighted values and weights independently, and then divides the sum of weighted values by the sum of weights at the server. Instead of implementing the summations by directly using the `tff.federated_sum` operator, the summation is delegated to two instances of `tff.aggregators.SumFactory`.\n",
        "\n",
        "Such structure makes it possible for the two default summations to be replaced by different factories, which realize the sum differently. For example, a `tff.aggregators.SecureSumFactory`, or a custom implementation of the `tff.aggregators.UnweightedAggregationFactory`. Conversely, time, `tff.aggregators.MeanFactory` can itself be an inner aggregation of another factory such as `tff.aggregators.clipping_factory`, if the values are to be clipped before averaging.\n",
        "\n",
        "See the previous [Tuning recommended aggregations for learning](tuning_recommended_aggregators.ipynb) tutorial for receommended uses of the composition mechanism using existing factories in the `tff.aggregators` module."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jPk4ZRQuiq0W"
      },
      "source": [
        "## Best practices by example\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1G8Q-SLXQKIv"
      },
      "source": [
        "We are going to illustrate the `tff.aggregators` concepts in detail by implementing a simple example task, and make it progressively more general. Another way to learn is to look at the implementation of existing factories."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sJukNEVporAx"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "\n",
        "import numpy as np\n",
        "import tensorflow as tf\n",
        "import tensorflow_federated as tff"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kuWnjpo9k_yG"
      },
      "source": [
        "Instead of summing `value`, the example task is to sum `value * 2.0` and then divide the sum by `2.0`. The aggregation result is thus mathematically equivalent to directly summing the `value`, and could be thought of as consisting of three parts: (1) scaling at clients (2) summing across clients (3) unscaling at server."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qpDHHDYMlC1D"
      },
      "source": [
        "NOTE: This task is not necessarily useful in practice. Nevertheless, it is helpful in explaining the underlying concepts."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CJnGAYmVlE-I"
      },
      "source": [
        "Following the design explained above, the logic will be implemented as a subclass of `tff.aggregators.UnweightedAggregationFactory`, which creates appropriate `tff.templates.AggregationProcess` when given a `value_type` to aggregate:"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N-3YKg8klWP0"
      },
      "source": [
        "### Minimal implementation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kIDclJTBlYDM"
      },
      "source": [
        "For the example task, the computations necessary are always the same, so there is no need for using state. It is thus empty, and represented as `tff.federated_value((), tff.SERVER)`. The same holds for measurements, for now.\n",
        "\n",
        "The minimal implementation of the task is thus as follows:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ShzVITY4Bv8X"
      },
      "outputs": [],
      "source": [
        "class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
        "\n",
        "  def create(self, value_type):\n",
        "    @tff.federated_computation()\n",
        "    def initialize_fn():\n",
        "      return tff.federated_value((), tff.SERVER)\n",
        "\n",
        "    @tff.federated_computation(initialize_fn.type_signature.result,\n",
        "                               tff.FederatedType(value_type, tff.CLIENTS))\n",
        "    def next_fn(state, value):\n",
        "      scaled_value = tff.federated_map(\n",
        "          tff.tf_computation(lambda x: x * 2.0), value)\n",
        "      summed_value = tff.federated_sum(scaled_value)\n",
        "      unscaled_value = tff.federated_map(\n",
        "          tff.tf_computation(lambda x: x / 2.0), summed_value)\n",
        "      measurements = tff.federated_value((), tff.SERVER)\n",
        "      return tff.templates.MeasuredProcessOutput(\n",
        "          state=state, result=unscaled_value, measurements=measurements)\n",
        "\n",
        "    return tff.templates.AggregationProcess(initialize_fn, next_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "XivHiQ38Y2Ph"
      },
      "source": [
        "Whether everything works as expected can be verified with the following code:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nd6qJL15oz1W"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Type signatures of the created aggregation process:\n",
            "  - initialize: ( -\u003e \u003c\u003e@SERVER)\n",
            "  - next: (\u003cstate=\u003c\u003e@SERVER,value={float32}@CLIENTS\u003e -\u003e \u003cstate=\u003c\u003e@SERVER,result=float32@SERVER,measurements=\u003c\u003e@SERVER\u003e)\n",
            "\n",
            "Aggregation result: 8.0  (expected 8.0)\n"
          ]
        }
      ],
      "source": [
        "client_data = [1.0, 2.0, 5.0]\n",
        "factory = ExampleTaskFactory()\n",
        "aggregation_process = factory.create(tff.TensorType(np.float32))\n",
        "print(f'Type signatures of the created aggregation process:\\n'\n",
        "      f'  - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
        "      f'  - next: {aggregation_process.next.type_signature}\\n')\n",
        "\n",
        "state = aggregation_process.initialize()\n",
        "output = aggregation_process.next(state, client_data)\n",
        "print(f'Aggregation result: {output.result}  (expected 8.0)')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WAdrM-JEpZjE"
      },
      "source": [
        "### Statefulness and measurements"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DLIu-aRQpo3k"
      },
      "source": [
        "Statefulness is broadly used in TFF to represent computations that are expected to be executed iteratively and change with each iteration. For example, the state of a learning computation contains the weights of the model being learned.\n",
        "\n",
        "To illustrate how to use state in an aggregation computation, we modify the example task. Instead of multiplying `value` by `2.0`, we multiply it by the iteration index - the number of times the aggregation has been executed.\n",
        "\n",
        "To do so, we need a way to keep track of the iteration index, which is achieved through the concept of state. In the `initialize_fn`, instead of creating an empty state, we initialize the state to be a scalar zero. Then, state can be used in the `next_fn` in three steps: (1) increment by `1.0`, (2) use to multiply `value`, and (3) return as the new updated state."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "efVS0FhEqELN"
      },
      "source": [
        "Once this is done, you may note: *But exactly the same code as above can be used to verify all works as expected. How do I know something has actually changed?*\n",
        "\n",
        "Good question! This is where the concept of measurements becomes useful. In general, measurements can report any value relevant to a single execution of the `next` function, which could be used for monitoring. In this case, it can be the `summed_value` from the previous example. That is, the value before the \"unscaling\" step, which should depend on the iteration index. *Again, this is not necessarily useful in practice, but illustrates the relevant mechanism.*\n",
        "\n",
        "The stateful answer to the task thus looks as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rADHGC0gqZhs"
      },
      "outputs": [],
      "source": [
        "class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
        "\n",
        "  def create(self, value_type):\n",
        "    @tff.federated_computation()\n",
        "    def initialize_fn():\n",
        "      return tff.federated_value(0.0, tff.SERVER)\n",
        "\n",
        "    @tff.federated_computation(initialize_fn.type_signature.result,\n",
        "                               tff.FederatedType(value_type, tff.CLIENTS))\n",
        "    def next_fn(state, value):\n",
        "      new_state = tff.federated_map(\n",
        "          tff.tf_computation(lambda x: x + 1.0), state)\n",
        "      state_at_clients = tff.federated_broadcast(new_state)\n",
        "      scaled_value = tff.federated_map(\n",
        "          tff.tf_computation(lambda x, y: x * y), (value, state_at_clients))\n",
        "      summed_value = tff.federated_sum(scaled_value)\n",
        "      unscaled_value = tff.federated_map(\n",
        "          tff.tf_computation(lambda x, y: x / y), (summed_value, new_state))\n",
        "      return tff.templates.MeasuredProcessOutput(\n",
        "          state=new_state, result=unscaled_value, measurements=summed_value)\n",
        "\n",
        "    return tff.templates.AggregationProcess(initialize_fn, next_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SIdlsupeqyTu"
      },
      "source": [
        "Note that the `state` that comes into `next_fn` as input is placed at server. In order to use it at clients, it first needs to be communicated, which is achieved using the `tff.federated_broadcast` operator.\n",
        "\n",
        "To verify all works as expected, we can now look at the reported `measurements`, which should be different with each round of execution, even if run with the same `client_data`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "W0E2jwodrO4R"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Type signatures of the created aggregation process:\n",
            "  - initialize: ( -\u003e float32@SERVER)\n",
            "  - next: (\u003cstate=float32@SERVER,value={float32}@CLIENTS\u003e -\u003e \u003cstate=float32@SERVER,result=float32@SERVER,measurements=float32@SERVER\u003e)\n",
            "\n",
            "| Round #1\n",
            "|       Aggregation result: 8.0   (expected 8.0)\n",
            "| Aggregation measurements: 8.0   (expected 8.0 * 1)\n",
            "\n",
            "| Round #2\n",
            "|       Aggregation result: 8.0   (expected 8.0)\n",
            "| Aggregation measurements: 16.0  (expected 8.0 * 2)\n",
            "\n",
            "| Round #3\n",
            "|       Aggregation result: 8.0   (expected 8.0)\n",
            "| Aggregation measurements: 24.0  (expected 8.0 * 3)\n"
          ]
        }
      ],
      "source": [
        "client_data = [1.0, 2.0, 5.0]\n",
        "factory = ExampleTaskFactory()\n",
        "aggregation_process = factory.create(tff.TensorType(np.float32))\n",
        "print(f'Type signatures of the created aggregation process:\\n'\n",
        "      f'  - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
        "      f'  - next: {aggregation_process.next.type_signature}\\n')\n",
        "\n",
        "state = aggregation_process.initialize()\n",
        "\n",
        "output = aggregation_process.next(state, client_data)\n",
        "print('| Round #1')\n",
        "print(f'|       Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| Aggregation measurements: {output.measurements}   (expected 8.0 * 1)')\n",
        "\n",
        "output = aggregation_process.next(output.state, client_data)\n",
        "print('\\n| Round #2')\n",
        "print(f'|       Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| Aggregation measurements: {output.measurements}  (expected 8.0 * 2)')\n",
        "\n",
        "output = aggregation_process.next(output.state, client_data)\n",
        "print('\\n| Round #3')\n",
        "print(f'|       Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| Aggregation measurements: {output.measurements}  (expected 8.0 * 3)')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wSruC1V9rxma"
      },
      "source": [
        "### Structured types"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yjoqb6Q_tcgl"
      },
      "source": [
        "The model weights of a model trained in federated learning are usually represented as a collection of tensors, rather than a single tensor. In TFF, this is represented as `tff.StructType` and generally useful aggregation factories need to be able to accept the structured types.\n",
        "\n",
        "However, in the above examples, we only worked with a `tff.TensorType` object. If we try to use the previous factory to create the aggregation process with a `tff.StructType([(np.float32, (2,)), (np.float32, (3,))])`, we get a strange error because TensorFlow will try to multiply a `tf.Tensor` and a `list`.\n",
        "\n",
        "The problem is that instead of multiplying the structure of tensors by a constant, we need to multiply *each tensor in the structure* by a constant. The usual solution to this problem is to use the `tf.nest` module inside of the created `tff.tf_computation`s.\n",
        "\n",
        "The version of the previous `ExampleTaskFactory` compatible with structured types thus looks as follows:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Hb4Wi-yUtgSs"
      },
      "outputs": [],
      "source": [
        "@tff.tf_computation()\n",
        "def scale(value, factor):\n",
        "  return tf.nest.map_structure(lambda x: x * factor, value)\n",
        "\n",
        "@tff.tf_computation()\n",
        "def unscale(value, factor):\n",
        "  return tf.nest.map_structure(lambda x: x / factor, value)\n",
        "\n",
        "@tff.tf_computation()\n",
        "def add_one(value):\n",
        "  return value + 1.0\n",
        "\n",
        "class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
        "\n",
        "  def create(self, value_type):\n",
        "    @tff.federated_computation()\n",
        "    def initialize_fn():\n",
        "      return tff.federated_value(0.0, tff.SERVER)\n",
        "\n",
        "    @tff.federated_computation(initialize_fn.type_signature.result,\n",
        "                               tff.FederatedType(value_type, tff.CLIENTS))\n",
        "    def next_fn(state, value):\n",
        "      new_state = tff.federated_map(add_one, state)\n",
        "      state_at_clients = tff.federated_broadcast(new_state)\n",
        "      scaled_value = tff.federated_map(scale, (value, state_at_clients))\n",
        "      summed_value = tff.federated_sum(scaled_value)\n",
        "      unscaled_value = tff.federated_map(unscale, (summed_value, new_state))\n",
        "      return tff.templates.MeasuredProcessOutput(\n",
        "          state=new_state, result=unscaled_value, measurements=summed_value)\n",
        "\n",
        "    return tff.templates.AggregationProcess(initialize_fn, next_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "x5-1gxv1uebz"
      },
      "source": [
        "This example highlights a pattern which may be useful to follow when structuring TFF code. When not dealing with very simple operations, the code becomes more legible when the `tff.tf_computation`s that will be used as building blocks inside a `tff.federated_computation` are created in a separate place. Inside of the `tff.federated_computation`, these building blocks are only connected using the intrinsic operators."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Uj9I-qKvu0_0"
      },
      "source": [
        "To verify it works as expected:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GLRZd0TFvAUR"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Type signatures of the created aggregation process:\n",
            "  - initialize: ( -\u003e float32@SERVER)\n",
            "  - next: (\u003cstate=float32@SERVER,value={\u003cfloat32[2],float32[3]\u003e}@CLIENTS\u003e -\u003e \u003cstate=float32@SERVER,result=\u003cfloat32[2],float32[3]\u003e@SERVER,measurements=\u003cfloat32[2],float32[3]\u003e@SERVER\u003e)\n",
            "\n",
            "Aggregation result: [[2. 3.], [6. 4. 0.]]\n",
            "          Expected: [[2. 3.], [6. 4. 0.]]\n"
          ]
        }
      ],
      "source": [
        "client_data = [[[1.0, 2.0], [3.0, 4.0, 5.0]],\n",
        "               [[1.0, 1.0], [3.0, 0.0, -5.0]]]\n",
        "factory = ExampleTaskFactory()\n",
        "aggregation_process = factory.create(\n",
        "    tff.to_type([(np.float32, (2,)), (np.float32, (3,))]))\n",
        "print(f'Type signatures of the created aggregation process:\\n'\n",
        "      f'  - initialize: {aggregation_process.initialize.type_signature}\\n'\n",
        "      f'  - next: {aggregation_process.next.type_signature}\\n')\n",
        "\n",
        "state = aggregation_process.initialize()\n",
        "output = aggregation_process.next(state, client_data)\n",
        "print(f'Aggregation result: [{output.result[0]}, {output.result[1]}]\\n'\n",
        "      f'          Expected: [[2. 3.], [6. 4. 0.]]')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BO8f11fgvl_4"
      },
      "source": [
        "### Inner aggregations"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YQWF8GWnx2U_"
      },
      "source": [
        "The final step is to optionally enable delegation of the actual aggregation to other factories, in order to allow easy composition of different aggregation techniques.\n",
        "\n",
        "This is achieved by creating an optional `inner_factory` argument in the constructor of our `ExampleTaskFactory`. If not specified, `tff.aggregators.SumFactory` is used, which applies the `tff.federated_sum` operator used directly in the previous section.\n",
        "\n",
        "When `create` is called, we can first call `create` of the `inner_factory` to create the inner aggregation process with the same `value_type`.\n",
        "\n",
        "The state of our process returned by `initialize_fn` is a composition of two parts: the state created by \"this\" process, and the state of the just created inner process.\n",
        "\n",
        "The implementation of the `next_fn` differs in that the actual aggregation is delegated to the `next` function of the inner process, and in how the final output is composed. The state is again composed of \"this\" and \"inner\" state, and measurements are composed in a similar manner as an `OrderedDict`.\n",
        "\n",
        "The following is an implementation of such pattern."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P0LxeArJx4xK"
      },
      "outputs": [],
      "source": [
        "@tff.tf_computation()\n",
        "def scale(value, factor):\n",
        "  return tf.nest.map_structure(lambda x: x * factor, value)\n",
        "\n",
        "@tff.tf_computation()\n",
        "def unscale(value, factor):\n",
        "  return tf.nest.map_structure(lambda x: x / factor, value)\n",
        "\n",
        "@tff.tf_computation()\n",
        "def add_one(value):\n",
        "  return value + 1.0\n",
        "\n",
        "class ExampleTaskFactory(tff.aggregators.UnweightedAggregationFactory):\n",
        "\n",
        "  def __init__(self, inner_factory=None):\n",
        "    if inner_factory is None:\n",
        "      inner_factory = tff.aggregators.SumFactory()\n",
        "    self._inner_factory = inner_factory\n",
        "\n",
        "  def create(self, value_type):\n",
        "    inner_process = self._inner_factory.create(value_type)\n",
        "\n",
        "    @tff.federated_computation()\n",
        "    def initialize_fn():\n",
        "      my_state = tff.federated_value(0.0, tff.SERVER)\n",
        "      inner_state = inner_process.initialize()\n",
        "      return tff.federated_zip((my_state, inner_state))\n",
        "\n",
        "    @tff.federated_computation(initialize_fn.type_signature.result,\n",
        "                               tff.FederatedType(value_type, tff.CLIENTS))\n",
        "    def next_fn(state, value):\n",
        "      my_state, inner_state = state\n",
        "      my_new_state = tff.federated_map(add_one, my_state)\n",
        "      my_state_at_clients = tff.federated_broadcast(my_new_state)\n",
        "      scaled_value = tff.federated_map(scale, (value, my_state_at_clients))\n",
        "\n",
        "      # Delegation to an inner factory, returning values placed at SERVER.\n",
        "      inner_output = inner_process.next(inner_state, scaled_value)\n",
        "\n",
        "      unscaled_value = tff.federated_map(unscale, (inner_output.result, my_new_state))\n",
        "\n",
        "      new_state = tff.federated_zip((my_new_state, inner_output.state))\n",
        "      measurements = tff.federated_zip(\n",
        "          collections.OrderedDict(\n",
        "              scaled_value=inner_output.result,\n",
        "              example_task=inner_output.measurements))\n",
        "\n",
        "      return tff.templates.MeasuredProcessOutput(\n",
        "          state=new_state, result=unscaled_value, measurements=measurements)\n",
        "\n",
        "    return tff.templates.AggregationProcess(initialize_fn, next_fn)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "36IkS3npyzrh"
      },
      "source": [
        "When delegating to the `inner_process.next` function, the return structure we get is a `tff.templates.MeasuredProcessOutput`, with the same three fields - `state`, `result` and `measurements`. When creating the overall return structure of the composed aggregation process, the `state` and `measurements` fields should be generally composed and returned together. In contrast, the `result` field corresponds to the value being aggregated and instead \"flows through\" the composed aggregation.\n",
        "\n",
        "The `state` object should be seen as an implementation detail of the factory, and thus the composition could be of any structure. However, `measurements` correspond to values to be reported to the user at some point. Therefore, we recommend to use `OrderedDict`, with composed naming such that it would be clear where in an composition does a reported metric comes from.\n",
        "\n",
        "Note also the use of the `tff.federated_zip` operator. The `state` object contolled by the created process should be a `tff.FederatedType`. If we had instead returned `(this_state, inner_state)` in the `initialize_fn`, its return type signature would be a `tff.StructType` containing a 2-tuple of `tff.FederatedType`s. The use of `tff.federated_zip` \"lifts\" the `tff.FederatedType` to the top level. This is similarly used in the `next_fn` when preparing the state and measurements to be returned."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "usIIplMl02zd"
      },
      "source": [
        "Finally, we can see how this can be used with the default inner aggregation:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NKCrTiVo1ACJ"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "| Round #1\n",
            "|           Aggregation result: 8.0   (expected 8.0)\n",
            "| measurements['scaled_value']: 8.0\n",
            "| measurements['example_task']: ()\n",
            "\n",
            "| Round #2\n",
            "|           Aggregation result: 8.0   (expected 8.0)\n",
            "| measurements['scaled_value']: 16.0\n",
            "| measurements['example_task']: ()\n"
          ]
        }
      ],
      "source": [
        "client_data = [1.0, 2.0, 5.0]\n",
        "factory = ExampleTaskFactory()\n",
        "aggregation_process = factory.create(tff.TensorType(np.float32))\n",
        "state = aggregation_process.initialize()\n",
        "\n",
        "output = aggregation_process.next(state, client_data)\n",
        "print('| Round #1')\n",
        "print(f'|           Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
        "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')\n",
        "\n",
        "output = aggregation_process.next(output.state, client_data)\n",
        "print('\\n| Round #2')\n",
        "print(f'|           Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
        "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B0cotpsl1n3c"
      },
      "source": [
        "... and with a different inner aggregation. For example, an `ExampleTaskFactory`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lhrsdsTH1ulS"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "| Round #1\n",
            "|           Aggregation result: 8.0   (expected 8.0)\n",
            "| measurements['scaled_value']: 8.0\n",
            "| measurements['example_task']: OrderedDict([('scaled_value', 8.0), ('example_task', ())])\n",
            "\n",
            "| Round #2\n",
            "|           Aggregation result: 8.0   (expected 8.0)\n",
            "| measurements['scaled_value']: 16.0\n",
            "| measurements['example_task']: OrderedDict([('scaled_value', 32.0), ('example_task', ())])\n"
          ]
        }
      ],
      "source": [
        "client_data = [1.0, 2.0, 5.0]\n",
        "# Note the inner delegation can be to any UnweightedAggregaionFactory.\n",
        "# In this case, each factory creates process that multiplies by the iteration\n",
        "# index (1, 2, 3, ...), thus their combination multiplies by (1, 4, 9, ...).\n",
        "factory = ExampleTaskFactory(ExampleTaskFactory())\n",
        "aggregation_process = factory.create(tff.TensorType(np.float32))\n",
        "state = aggregation_process.initialize()\n",
        "\n",
        "output = aggregation_process.next(state, client_data)\n",
        "print('| Round #1')\n",
        "print(f'|           Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
        "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')\n",
        "\n",
        "output = aggregation_process.next(output.state, client_data)\n",
        "print('\\n| Round #2')\n",
        "print(f'|           Aggregation result: {output.result}   (expected 8.0)')\n",
        "print(f'| measurements[\\'scaled_value\\']: {output.measurements[\"scaled_value\"]}')\n",
        "print(f'| measurements[\\'example_task\\']: {output.measurements[\"example_task\"]}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Br9GTYu40GKY"
      },
      "source": [
        "## Summary\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "O-Wpu5hG0Jdl"
      },
      "source": [
        "In this tutorial, we explained the best practices to follow in order to create a general-purpose aggregation building block, represented as an aggregation factory. The generality comes through the design intent in two ways:\n",
        "\n",
        "1. *Parameterized computations.* Aggregation is an independent building block that can be plugged into other TFF modules designed to work with `tff.aggregators` to parameterize their necessary aggregation, such as `tff.learning.algorithms.build_weighted_fed_avg`.\n",
        "1. *Aggregation composition.* An aggregation building block can be composed with other aggregation building blocks to create more complex composite aggregations."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "a930wM_fqUNH"
      ],
      "name": "custom_aggregators.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}