site/en-snapshot/guide/random_numbers.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "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": "qFdPvlXBOdUN"
      },
      "source": [
        "# Random number generation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://www.tensorflow.org/guide/random_numbers\"><img src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" />View on TensorFlow.org</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs/blob/master/site/en/guide/random_numbers.ipynb\"><img src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" />Run in Google Colab</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a target=\"_blank\" href=\"https://github.com/tensorflow/docs/blob/master/site/en/guide/random_numbers.ipynb\"><img src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" />View source on GitHub</a>\n",
        "  </td>\n",
        "  <td>\n",
        "    <a href=\"https://storage.googleapis.com/tensorflow_docs/docs/site/en/guide/random_numbers.ipynb\"><img src=\"https://www.tensorflow.org/images/download_logo_32px.png\" />Download notebook</a>\n",
        "  </td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BlGY1iiph_C2"
      },
      "source": [
        "TensorFlow provides a set of pseudo-random number generators (RNG), in the `tf.random` module. This document describes how you can control the random number generators, and how these generators interact with other tensorflow sub-systems. \n",
        "\n",
        "Note: The random numbers are not guaranteed to be consistent across TensorFlow versions. See: [Version Compatibility](https://www.tensorflow.org/guide/versions#what_is_not_covered)\n",
        "\n",
        "TensorFlow provides two approaches for controlling the random number generation process:\n",
        "\n",
        "1. Through the explicit use of `tf.random.Generator` objects. Each such object maintains a state (in `tf.Variable`) that will be changed after each number generation.\n",
        "\n",
        "2. Through the purely-functional stateless random functions like `tf.random.stateless_uniform`. Calling these functions with the same arguments (which include the seed) and on the same device will always produce the same results.\n",
        "\n",
        "Warning: The old RNGs from TF 1.x such as `tf.random.uniform` and `tf.random.normal` are not yet deprecated but strongly discouraged."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zIGh9faCOp6x"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ECDrttf0s8Nu"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "# Creates some virtual devices (cpu:0, cpu:1, etc.) for using distribution strategy\n",
        "physical_devices = tf.config.list_physical_devices(\"CPU\")\n",
        "tf.config.experimental.set_virtual_device_configuration(\n",
        "    physical_devices[0], [\n",
        "        tf.config.experimental.VirtualDeviceConfiguration(),\n",
        "        tf.config.experimental.VirtualDeviceConfiguration(),\n",
        "        tf.config.experimental.VirtualDeviceConfiguration()\n",
        "    ])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eqMlrUsVu2Ai"
      },
      "source": [
        "## The `tf.random.Generator` class\n",
        "\n",
        "The `tf.random.Generator` class is used in cases where you want each RNG call to produce different results. It maintains an internal state (managed by a `tf.Variable` object) which will be updated every time random numbers are generated. Because the state is managed by `tf.Variable`, it enjoys all facilities provided by `tf.Variable` such as easy checkpointing, automatic control-dependency and thread safety.\n",
        "\n",
        "You can get a `tf.random.Generator` by manually creating an object of the class or call `tf.random.get_global_generator()` to get the default global generator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7yU1E3JvxOQD"
      },
      "outputs": [],
      "source": [
        "g1 = tf.random.Generator.from_seed(1)\n",
        "print(g1.normal(shape=[2, 3]))\n",
        "g2 = tf.random.get_global_generator()\n",
        "print(g2.normal(shape=[2, 3]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QmRCeAvTxulW"
      },
      "source": [
        "There are multiple ways to create a generator object. The easiest is `Generator.from_seed`, as shown above, that creates a generator from a seed. A seed is any non-negative integer. `from_seed` also takes an optional argument `alg` which is the RNG algorithm that will be used by this generator:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kISbOE4Xfjhv"
      },
      "outputs": [],
      "source": [
        "g1 = tf.random.Generator.from_seed(1, alg='philox')\n",
        "print(g1.normal(shape=[2, 3]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_mCRaN7dfd8j"
      },
      "source": [
        "See the *Algorithms* section below for more information about it.\n",
        "\n",
        "Another way to create a generator is with `Generator.from_non_deterministic_state`. A generator created this way will start from a non-deterministic state, depending on e.g. time and OS."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "gxPLCLsz00qY"
      },
      "outputs": [],
      "source": [
        "g = tf.random.Generator.from_non_deterministic_state()\n",
        "print(g.normal(shape=[2, 3]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zSAp2BMj1JZ6"
      },
      "source": [
        "There are yet other ways to create generators, such as from explicit states, which are not covered by this guide.\n",
        "\n",
        "When using `tf.random.get_global_generator` to get the global generator, you need to be careful about device placement. The global generator is created (from a non-deterministic state) at the first time `tf.random.get_global_generator` is called, and placed on the default device at that call. So, for example, if the first site you call `tf.random.get_global_generator` is within a `tf.device(\"gpu\")` scope, the global generator will be placed on the GPU, and using the global generator later on from the CPU will incur a GPU-to-CPU copy.\n",
        "\n",
        "There is also a function `tf.random.set_global_generator` for replacing the global generator with another generator object. This function should be used with caution though, because the old global generator may have been captured by a `tf.function` (as a weak reference), and replacing it will cause it to be garbage collected, breaking the `tf.function`. A better way to reset the global generator is to use one of the \"reset\" functions such as `Generator.reset_from_seed`, which won't create new generator objects."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "324S5bpd9HRg"
      },
      "outputs": [],
      "source": [
        "g = tf.random.Generator.from_seed(1)\n",
        "print(g.normal([]))\n",
        "print(g.normal([]))\n",
        "g.reset_from_seed(1)\n",
        "print(g.normal([]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "z9H0wuvp9VwH"
      },
      "source": [
        "### Creating independent random-number streams\n",
        "\n",
        "In many applications one needs multiple independent random-number streams, independent in the sense that they won't overlap and won't have any statistically detectable correlations. This is achieved by using `Generator.split` to create multiple generators that are guaranteed to be independent of each other (i.e. generating independent streams)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vg5_KN18OZjo"
      },
      "outputs": [],
      "source": [
        "g = tf.random.Generator.from_seed(1)\n",
        "print(g.normal([]))\n",
        "new_gs = g.split(3)\n",
        "for new_g in new_gs:\n",
        "  print(new_g.normal([]))\n",
        "print(g.normal([]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dqOaGVzKOsRJ"
      },
      "source": [
        "`split` will change the state of the generator on which it is called (`g` in the above example), similar to an RNG method such as `normal`. In addition to being independent of each other, the new generators (`new_gs`) are also guaranteed to be independent of the old one (`g`).\n",
        "\n",
        "Spawning new generators is also useful when you want to make sure the generator you use is on the same device as other computations, to avoid the overhead of cross-device copy. For example: "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5jSnJBlUQzF3"
      },
      "outputs": [],
      "source": [
        "with tf.device(\"cpu\"):  # change \"cpu\" to the device you want\n",
        "  g = tf.random.get_global_generator().split(1)[0]  \n",
        "  print(g.normal([]))  # use of g won't cause cross-device copy, unlike the global generator"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sCxbccYMRdd4"
      },
      "source": [
        "Note: In theory, you can use constructors such as `from_seed` instead of `split` here to obtain a new generator, but by doing so you lose the guarantee that the new generator is independent of the global generator. You will also run the risk that you may accidentally create two generators with the same seed or with seeds that lead to overlapping random-number streams.\n",
        "\n",
        "You can do splitting recursively, calling `split` on split generators. There are no limits (barring integer overflow) on the depth of recursions."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8JUgnQM_O0lg"
      },
      "source": [
        "### Interaction with `tf.function`\n",
        "\n",
        "`tf.random.Generator` obeys the same rules as `tf.Variable` when used with `tf.function`. This includes three aspects."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jnSjhY6WM-J8"
      },
      "source": [
        "#### Creating generators outside `tf.function` \n",
        "\n",
        "`tf.function` can use a generator created outside of it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "a5EEy0E2UHMw"
      },
      "outputs": [],
      "source": [
        "g = tf.random.Generator.from_seed(1)\n",
        "@tf.function\n",
        "def foo():\n",
        "  return g.normal([])\n",
        "print(foo())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L_8kC7kbO5uu"
      },
      "source": [
        "The user needs to make sure that the generator object is still alive (not garbage-collected) when the function is called."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PwIrBv_zUYwI"
      },
      "source": [
        "#### Creating generators inside `tf.function` \n",
        "\n",
        "Creation of generators inside a `tf.function` can only happened during the first run of the function. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3JzpUvqJU4MW"
      },
      "outputs": [],
      "source": [
        "g = None\n",
        "@tf.function\n",
        "def foo():\n",
        "  global g\n",
        "  if g is None:\n",
        "    g = tf.random.Generator.from_seed(1)\n",
        "  return g.normal([])\n",
        "print(foo())\n",
        "print(foo())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UaTVnOhHVM9a"
      },
      "source": [
        "#### Passing generators as arguments to `tf.function`\n",
        "\n",
        "When used as an argument to a `tf.function`, different generator objects will cause retracing of the `tf.function`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "DeR9kvt0V-ad"
      },
      "outputs": [],
      "source": [
        "num_traces = 0\n",
        "@tf.function\n",
        "def foo(g):\n",
        "  global num_traces\n",
        "  num_traces += 1\n",
        "  return g.normal([])\n",
        "foo(tf.random.Generator.from_seed(1))\n",
        "foo(tf.random.Generator.from_seed(2))\n",
        "print(num_traces)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "E0RxllJzkGfo"
      },
      "source": [
        "Note that this retracing behavior is consistent with `tf.Variable`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "oWD2f_qxkSe7"
      },
      "outputs": [],
      "source": [
        "num_traces = 0\n",
        "@tf.function\n",
        "def foo(v):\n",
        "  global num_traces\n",
        "  num_traces += 1\n",
        "  return v.read_value()\n",
        "foo(tf.Variable(1))\n",
        "foo(tf.Variable(2))\n",
        "print(num_traces)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fxcS6IY8WZuh"
      },
      "source": [
        "### Interaction with distribution strategies\n",
        "\n",
        "There are two ways in which `Generator` interacts with distribution strategies."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GyZv9QJkZfkQ"
      },
      "source": [
        "#### Creating generators outside distribution strategies\n",
        "\n",
        "If a generator is created outside strategy scopes, all replicas’ access to the generator will be serialized, and hence the replicas will get different random numbers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HX_beT9SZWMp"
      },
      "outputs": [],
      "source": [
        "g = tf.random.Generator.from_seed(1)\n",
        "strat = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\"])\n",
        "with strat.scope():\n",
        "  def f():\n",
        "    print(g.normal([]))\n",
        "  results = strat.run(f)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ydYQbUqLPAgH"
      },
      "source": [
        "Note that this usage may have performance issues because the generator's device is different from the replicas."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Yal4LbBKbAeN"
      },
      "source": [
        "#### Creating generators inside distribution strategies\n",
        "\n",
        "If a generator is created inside a strategy scope, each replica will get a different and independent stream of random numbers."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5SeUu7IFmTyQ"
      },
      "outputs": [],
      "source": [
        "strat = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\"])\n",
        "with strat.scope():\n",
        "  g = tf.random.Generator.from_seed(1)\n",
        "  print(strat.run(lambda: g.normal([])))\n",
        "  print(strat.run(lambda: g.normal([])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PFBlrOudfu9u"
      },
      "source": [
        "Note: Currently `tf.random.Generator` doesn't provide an option to let different replicas get identical (instead of different) streams (which is technically not hard). If you have a use case for this feature, please let the TensorFlow developers know.\n",
        "\n",
        "If the generator is seeded (e.g. created by `Generator.from_seed`), the random numbers are determined by the seed, even though different replicas get different and uncorrelated numbers. One can think of a random number generated on a replica as a hash of the replica ID and a \"primary\" random number that is common to all replicas. Hence, the whole system is still deterministic.\n",
        "\n",
        "`tf.random.Generator` can also be created inside `Strategy.run`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nlQXi5Msb1Wu"
      },
      "outputs": [],
      "source": [
        "strat = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\"])\n",
        "with strat.scope():\n",
        "  def f():\n",
        "    g = tf.random.Generator.from_seed(1)\n",
        "    a = g.normal([])\n",
        "    b = g.normal([])\n",
        "    return tf.stack([a, b])\n",
        "  print(strat.run(f))\n",
        "  print(strat.run(f))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4Sv-aiaOmrOr"
      },
      "source": [
        "We no longer recommend passing `tf.random.Generator` as arguments to `Strategy.run`, because `Strategy.run` generally expects the arguments to be tensors, not generators."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8RbM4vabtiWM"
      },
      "source": [
        "### Saving generators\n",
        "\n",
        "Generally for saving or serializing you can handle a `tf.random.Generator` the same way you would handle a `tf.Variable` or a `tf.Module` (or its subclasses). In TF there are two mechanisms for serialization: [Checkpoint](https://www.tensorflow.org/guide/checkpoint) and [SavedModel](https://www.tensorflow.org/guide/saved_model)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PDtySQDotWQc"
      },
      "source": [
        "#### Checkpoint\n",
        "\n",
        "Generators can be freely saved and restored using `tf.train.Checkpoint`. The random-number stream from the restoring point will be the same as that from the saving point. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uB_bDSbzpbne"
      },
      "outputs": [],
      "source": [
        "filename = \"./checkpoint\"\n",
        "g = tf.random.Generator.from_seed(1)\n",
        "cp = tf.train.Checkpoint(generator=g)\n",
        "print(g.normal([]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bKKtRWeIkIjX"
      },
      "outputs": [],
      "source": [
        "cp.write(filename)\n",
        "print(\"RNG stream from saving point:\")\n",
        "print(g.normal([]))\n",
        "print(g.normal([]))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-cIHcHwRkQp3"
      },
      "outputs": [],
      "source": [
        "cp.restore(filename)\n",
        "print(\"RNG stream from restoring point:\")\n",
        "print(g.normal([]))\n",
        "print(g.normal([]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "A-OeUUQEJ37X"
      },
      "source": [
        "You can also save and restore within a distribution strategy:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3aI6TQ2lq28w"
      },
      "outputs": [],
      "source": [
        "filename = \"./checkpoint\"\n",
        "strat = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\"])\n",
        "with strat.scope():\n",
        "  g = tf.random.Generator.from_seed(1)\n",
        "  cp = tf.train.Checkpoint(my_generator=g)\n",
        "  print(strat.run(lambda: g.normal([])))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kTZcdaMwkvJI"
      },
      "outputs": [],
      "source": [
        "with strat.scope():\n",
        "  cp.write(filename)\n",
        "  print(\"RNG stream from saving point:\")\n",
        "  print(strat.run(lambda: g.normal([])))\n",
        "  print(strat.run(lambda: g.normal([])))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nizFA5IrkzN1"
      },
      "outputs": [],
      "source": [
        "with strat.scope():\n",
        "  cp.restore(filename)\n",
        "  print(\"RNG stream from restoring point:\")\n",
        "  print(strat.run(lambda: g.normal([])))\n",
        "  print(strat.run(lambda: g.normal([])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z2rsPfp9J6JA"
      },
      "source": [
        "You should make sure that the replicas don't diverge in their RNG call history (e.g. one replica makes one RNG call while another makes two RNG calls) before saving. Otherwise, their internal RNG states will diverge and `tf.train.Checkpoint` (which only saves the first replica's state) won't properly restore all the replicas.\n",
        "\n",
        "You can also restore a saved checkpoint to a different distribution strategy with a different number of replicas. Because a `tf.random.Generator` object created in a strategy can only be used in the same strategy, to restore to a different strategy, you have to create a new `tf.random.Generator` in the target strategy and a new `tf.train.Checkpoint` for it, as shown in this example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zgoFRf59-IvW"
      },
      "outputs": [],
      "source": [
        "filename = \"./checkpoint\"\n",
        "strat1 = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\"])\n",
        "with strat1.scope():\n",
        "  g1 = tf.random.Generator.from_seed(1)\n",
        "  cp1 = tf.train.Checkpoint(my_generator=g1)\n",
        "  print(strat1.run(lambda: g1.normal([])))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Lu79ETxMlDpO"
      },
      "outputs": [],
      "source": [
        "with strat1.scope():\n",
        "  cp1.write(filename)\n",
        "  print(\"RNG stream from saving point:\")\n",
        "  print(strat1.run(lambda: g1.normal([])))\n",
        "  print(strat1.run(lambda: g1.normal([])))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VYoRFUjklKOk"
      },
      "outputs": [],
      "source": [
        "strat2 = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\", \"cpu:2\"])\n",
        "with strat2.scope():\n",
        "  g2 = tf.random.Generator.from_seed(1)\n",
        "  cp2 = tf.train.Checkpoint(my_generator=g2)\n",
        "  cp2.restore(filename)\n",
        "  print(\"RNG stream from restoring point:\")\n",
        "  print(strat2.run(lambda: g2.normal([])))\n",
        "  print(strat2.run(lambda: g2.normal([])))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kMltUKbANqgl"
      },
      "source": [
        "Although `g1` and `cp1` are different objects from `g2` and `cp2`, they are linked via the common checkpoint file `filename` and object name `my_generator`. Overlapping replicas between strategies (e.g. `cpu:0` and `cpu:1` above) will have their RNG streams properly restored like in previous examples. This guarantee doesn't cover the case when a generator is saved in a strategy scope and restored outside of any strategy scope or vice versa, because a device outside strategies is treated as different from any replica in a strategy."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "w9dqrp1LnTaJ"
      },
      "source": [
        "#### SavedModel\n",
        "\n",
        "`tf.random.Generator` can be saved to a SavedModel. The generator can be created within a strategy scope. The saving can also happen within a strategy scope. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0AKO5SnUtyqx"
      },
      "outputs": [],
      "source": [
        "filename = \"./saved_model\"\n",
        "\n",
        "class MyModule(tf.Module):\n",
        "\n",
        "  def __init__(self):\n",
        "    super(MyModule, self).__init__()\n",
        "    self.g = tf.random.Generator.from_seed(0)\n",
        "\n",
        "  @tf.function\n",
        "  def __call__(self):\n",
        "    return self.g.normal([])\n",
        "\n",
        "  @tf.function\n",
        "  def state(self):\n",
        "    return self.g.state\n",
        "\n",
        "strat = tf.distribute.MirroredStrategy(devices=[\"cpu:0\", \"cpu:1\"])\n",
        "with strat.scope():\n",
        "  m = MyModule()\n",
        "  print(strat.run(m))\n",
        "  print(\"state:\", m.state())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jg2148hulfLB"
      },
      "outputs": [],
      "source": [
        "with strat.scope():\n",
        "  tf.saved_model.save(m, filename)\n",
        "  print(\"RNG stream from saving point:\")\n",
        "  print(strat.run(m))\n",
        "  print(\"state:\", m.state())\n",
        "  print(strat.run(m))\n",
        "  print(\"state:\", m.state())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "93AgVyzOllG7"
      },
      "outputs": [],
      "source": [
        "imported = tf.saved_model.load(filename)\n",
        "print(\"RNG stream from loading point:\")\n",
        "print(\"state:\", imported.state())\n",
        "print(imported())\n",
        "print(\"state:\", imported.state())\n",
        "print(imported())\n",
        "print(\"state:\", imported.state())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sbb23j3pZNNq"
      },
      "source": [
        "Loading a SavedModel containing `tf.random.Generator` into a distribution strategy is not recommended because the replicas will all generate the same random-number stream (which is because replica ID is frozen in SavedModel's graph). \n",
        "\n",
        "Loading a distributed `tf.random.Generator` (a generator created within a distribution strategy) into a non-strategy environment, like the above example, also has a caveat. The RNG state will be properly restored, but the random numbers generated will be different from the original generator in its strategy (again because a device outside strategies is treated as different from any replica in a strategy)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "73an1POpsi6V"
      },
      "source": [
        "## Stateless RNGs\n",
        "\n",
        "Usage of stateless RNGs is simple. Since they are just pure functions, there is no state or side effect involved."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0-aOOA3gasn_"
      },
      "outputs": [],
      "source": [
        "print(tf.random.stateless_normal(shape=[2, 3], seed=[1, 2]))\n",
        "print(tf.random.stateless_normal(shape=[2, 3], seed=[1, 2]))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2O_D-RAFNH2Q"
      },
      "source": [
        "Every stateless RNG requires a `seed` argument, which needs to be an integer Tensor of shape `[2]`. The results of the op are fully determined by this seed.\n",
        "\n",
        "The RNG algorithm used by stateless RNGs is device-dependent, meaning the same op running on a different device may produce different outputs."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4BvGkPnaOUPF"
      },
      "source": [
        "## Algorithms"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "58-8kvR4pRwO"
      },
      "source": [
        "### General\n",
        "\n",
        "Both the `tf.random.Generator` class and the `stateless` functions support the Philox algorithm (written as `\"philox\"` or `tf.random.Algorithm.PHILOX`) on all devices.\n",
        "\n",
        "Different devices will generate the same integer numbers, if using the same algorithm and starting from the same state. They will also generate \"almost the same\" float-point numbers, though there may be small numerical discrepancies caused by the different ways the devices carry out the float-point computation (e.g. reduction order)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WETA04F1OYPL"
      },
      "source": [
        "### XLA devices\n",
        "\n",
        "On XLA-driven devices (such as TPU, and also CPU/GPU when XLA is enabled) the ThreeFry algorithm (written as `\"threefry\"` or `tf.random.Algorithm.THREEFRY`) is also supported. This algorithm is fast on TPU but slow on CPU/GPU compared to Philox. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "c04JkebCPTPu"
      },
      "source": [
        "See paper ['Parallel Random Numbers: As Easy as 1, 2, 3'](https://www.thesalmons.org/john/random123/papers/random123sc11.pdf) for more details about these algorithms."
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "random_numbers.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}