Repository: google/flax

Issue:
Instance Normalization
Could you implement Instance Normalization, please?

Discussion:
codecov-commenter: ## [Codecov](https://app.codecov.io/gh/google/flax/pull/3652?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) Report
Attention: `4 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`1d23425`)](https://app.codecov.io/gh/google/flax/commit/1d2342513ad547adb8a5b0b717afe6db840ace4c?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 56.25% compared to head [(`5a49711`)](https://app.codecov.io/gh/google/flax/pull/3652?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) 56.28%.

| [Files](https://app.codecov.io/gh/google/flax/pull/3652?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) | Patch % | Lines |
|---|---|---|
| [flax/linen/normalization.py](https://app.codecov.io/gh/google/flax/pull/3652?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google#diff-ZmxheC9saW5lbi9ub3JtYWxpemF0aW9uLnB5) | 77.77% | [4 Missing :warning: ](https://app.codecov.io/gh/google/flax/pull/3652?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #3652      +/-   ##
==========================================
+ Coverage   56.25%   56.28%   +0.03%     
==========================================
  Files         101      101              
  Lines       12177    12195      +18     
==========================================
+ Hits         6850     6864      +14     
- Misses       5327     5331       +4     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/google/flax/pull/3652?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=google).


Patch:
diff --git a/docs/api_reference/flax.linen/layers.rst b/docs/api_reference/flax.linen/layers.rst
index d554deef..f732675b 100644
--- a/docs/api_reference/flax.linen/layers.rst
+++ b/docs/api_reference/flax.linen/layers.rst
@@ -56,6 +56,10 @@ Normalization
   :module: flax.linen
   :class: RMSNorm
 
+.. flax_module::
+  :module: flax.linen
+  :class: InstanceNorm
+
 .. flax_module::
   :module: flax.linen
   :class: SpectralNorm
diff --git a/docs/guides/flax_fundamentals/index.rst b/docs/guides/flax_fundamentals/index.rst
index 04cf5e9b..662cdb91 100644
--- a/docs/guides/flax_fundamentals/index.rst
+++ b/docs/guides/flax_fundamentals/index.rst
@@ -8,5 +8,4 @@ Flax fundamentals
    flax_basics
    state_params
    setup_or_nncompact
-   arguments
-   rng_guide
\ No newline at end of file
+   arguments
\ No newline at end of file
diff --git a/docs/guides/flax_fundamentals/rng_guide.ipynb b/docs/guides/flax_fundamentals/rng_guide.ipynb
deleted file mode 100644
index 0cecc3cb..00000000
--- a/docs/guides/flax_fundamentals/rng_guide.ipynb
+++ /dev/null
@@ -1,1847 +0,0 @@
-{
- "cells": [
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "# Randomness and PRNGs in Flax"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "In this guide, you will learn how Flax uses [JAX's explicit pseudorandom number generator (PRNG) keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#jax-prng) to emulate randomness, and adds some additional features to make it easier for users to thread PRNG keys through different Flax `Module`s.\n",
-    "\n",
-    "If you are new to JAX PRNG keys or need a refresher, check out:\n",
-    "- [JAX 101: PRNGs in JAX](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)\n",
-    "- [🔪 JAX - The Sharp Bits 🔪: Random Numbers](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers)."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Setup"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Install or upgrade Flax, and then import some necessary dependencies.\n",
-    "\n",
-    "**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don’t need this if you are already using a multi-device Google Cloud TPU environment, for example, on Google Cloud or in a Kaggle VM with a TPU."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 42,
-   "metadata": {
-    "outputId": "bb13d0ba-f908-4746-e4d3-f5916e0670ff",
-    "tags": [
-     "skip-execution"
-    ]
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "/bin/sh: line 1: pip: command not found\n"
-     ]
-    }
-   ],
-   "source": [
-    "!pip install -q flax"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 43,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import os\n",
-    "os.environ[\"XLA_FLAGS\"] = '--xla_force_host_platform_device_count=8'"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 44,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import flax, flax.linen as nn\n",
-    "import jax, jax.numpy as jnp\n",
-    "from jax.sharding import Mesh, PartitionSpec, NamedSharding\n",
-    "from jax.experimental import mesh_utils\n",
-    "from jax.experimental.shard_map import shard_map\n",
-    "\n",
-    "import hashlib"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 45,
-   "metadata": {
-    "outputId": "ab5e0d3b-2d51-44ee-b823-d152ad1a10b2"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[CpuDevice(id=0),\n",
-       " CpuDevice(id=1),\n",
-       " CpuDevice(id=2),\n",
-       " CpuDevice(id=3),\n",
-       " CpuDevice(id=4),\n",
-       " CpuDevice(id=5),\n",
-       " CpuDevice(id=6),\n",
-       " CpuDevice(id=7)]"
-      ]
-     },
-     "execution_count": 45,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "jax.devices()"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Set the JAX config variable `jax_threefry_partitionable` to `True`. This will be the default value in the future and makes the PRNG more efficiently auto-parallelizable under `jax.jit`. Refer to [JAX discussion](https://github.com/google/jax/discussions/18480) for more details."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 46,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "jax.config.update('jax_threefry_partitionable', True)\n",
-    "assert jax.config.jax_threefry_partitionable == True\n",
-    "assert jax.config.jax_default_prng_impl == 'threefry2x32'"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Receiving, manipulating and creating PRNG keys with `Module.make_rng`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The primary method Flax uses to receive, manipulate and create PRNG keys is via the `Module` method [`self.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng). It is a method that accepts a string name that represents an \"RNG stream\". Each RNG stream has an initial starting seed PRNG key, which the user passes in as a dictionary argument (i.e. into an [`.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) or [`.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) function), and the starting seed is used by `self.make_rng` to generate more PRNG keys for that stream.\n",
-    "\n",
-    "Note that this method can only be called with bounded modules (see [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html#top-level-modules))."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 47,
-   "metadata": {
-    "outputId": "8304fbc2-ab0a-42f0-9d4c-c88a74012e83"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[1915779057 2258748098]\n"
-     ]
-    }
-   ],
-   "source": [
-    "class RNGModule(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self):\n",
-    "    print(self.make_rng('rng_stream'))\n",
-    "    print(self.make_rng('rng_stream'))\n",
-    "    print(self.make_rng('rng_stream'))\n",
-    "\n",
-    "rng_module = RNGModule()\n",
-    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Now if we use a different starting seed PRNG key, we will generate different values (as intended)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 48,
-   "metadata": {
-    "outputId": "1cfd9632-11cf-43f4-b922-63038705a195"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[1830439201 4095528436]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[3737706588 1614077470]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[2940838374 2782395343]\n"
-     ]
-    }
-   ],
-   "source": [
-    "variables = rng_module.init({'rng_stream': jax.random.key(1)})"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Calling `self.make_rng` for one stream will not affect the random values generated from another stream; i.e. the call order doesn't matter."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 49,
-   "metadata": {
-    "outputId": "60dde8ff-08a9-4743-c7ec-c997aa0537b6"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
-      "[1830439201 4095528436]\n",
-      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
-      "[3737706588 1614077470]\n",
-      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1915779057 2258748098]\n",
-      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
-      "[2940838374 2782395343]\n"
-     ]
-    }
-   ],
-   "source": [
-    "class RNGModuleTwoStreams(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self):\n",
-    "    # same value as first code snippet above\n",
-    "    print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n",
-    "    # same value as second code snippet above\n",
-    "    print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n",
-    "    # same value as first code snippet above\n",
-    "    print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n",
-    "    # same value as second code snippet above\n",
-    "    print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n",
-    "    # same value as first code snippet above\n",
-    "    print(f\"rng_stream1: {self.make_rng('rng_stream1')}\")\n",
-    "    # same value as second code snippet above\n",
-    "    print(f\"rng_stream2: {self.make_rng('rng_stream2')}\")\n",
-    "\n",
-    "rng_module_two_streams = RNGModuleTwoStreams()\n",
-    "variables = rng_module_two_streams.init(\n",
-    "  {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(1)}\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Providing the same seed PRNG key will result in the same values being generated (provided that the same operations are used for those keys)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 50,
-   "metadata": {
-    "outputId": "12e8c3c5-a32c-46ab-ba71-98cc90aaed70"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "rng_stream1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1915779057 2258748098]\n",
-      "rng_stream2: Array((), dtype=key<fry>) overlaying:\n",
-      "[1915779057 2258748098]\n"
-     ]
-    }
-   ],
-   "source": [
-    "variables = rng_module_two_streams.init(\n",
-    "  {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(0)}\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### How `self.make_rng` works under the hood"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This is what happens when `self.make_rng` (`flax.linen.Module.make_rng`) is called:\n",
-    "* The following data is collected:\n",
-    "  * The path of the `Module` as provided by `self.scope.path` (the top-level root module has an empty path `()`).\n",
-    "  * The `self.make_rng` call count. That is, the number of times `self.make_rng` has been called for this specific stream (including this call).\n",
-    "    * **Note:** Each sub-`Module` will have its own individual call count that's separate from other `Module`s. For example, a `Module` that has called `self.make_rng('params')` twice and contains a sub-`Module` that has called `self.make_rng('params')` once, will have a call count of 2 and 1 for each of the RNG stream `'params'`, respectively.\n",
-    "* The data is bundled into a tuple and fed into a hash function and produces an integer.\n",
-    "* The generated integer is folded into the RNG stream's starting seed PRNG key to generate a new, unique PRNG key."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Below is a slightly simplified version of the hash function that Flax uses for `self.make_rng`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 51,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def produce_hash(data):\n",
-    "  m = hashlib.sha1()\n",
-    "  for x in data:\n",
-    "    if isinstance(x, str):\n",
-    "      m.update(x.encode('utf-8'))\n",
-    "    elif isinstance(x, int):\n",
-    "      m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))\n",
-    "    else:\n",
-    "      raise ValueError(f'Expected int or string, got: {x}')\n",
-    "  d = m.digest()\n",
-    "  hash_int = int.from_bytes(d[:4], byteorder='big')\n",
-    "  return hash_int"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "And now you can manually reproduce the PRNG keys generated from `self.make_rng`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 52,
-   "metadata": {
-    "outputId": "6bf352c1-086a-4c7a-8ab9-92abde614270"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[1428664606 3351135085]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[3456700291 3873160899]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[2411773124 4124888837]\n"
-     ]
-    }
-   ],
-   "source": [
-    "stream_seed = jax.random.key(0)\n",
-    "for call_count in range(1, 4):\n",
-    "  hash_int = produce_hash(data=(call_count,))\n",
-    "  print(jax.random.fold_in(stream_seed, jnp.uint32(hash_int)))"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 53,
-   "metadata": {
-    "outputId": "3809fddc-7319-4d55-ccc5-7b7c78160ea9"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "Array((), dtype=key<fry>) overlaying:\n",
-      "[1915779057 2258748098]\n"
-     ]
-    }
-   ],
-   "source": [
-    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Sub-`Module`s and `self.make_rng`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This section explores how `self.make_rng` (`flax.linen.Module.make_rng`) behaves with sub-`Module`s.\n",
-    "\n",
-    "Consider the following example:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 54,
-   "metadata": {
-    "outputId": "cc88c770-80ff-4093-e37f-6dc10c24006f"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[3619592043  626287670]\n",
-      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[965377860 480622172]\n",
-      "RNGSubSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1015683150 3648653849]\n",
-      "RNGSubSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[3694284925 2979568433]\n"
-     ]
-    }
-   ],
-   "source": [
-    "class RNGSubSubModule(nn.Module):\n",
-    "  def __call__(self):\n",
-    "    print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n",
-    "    print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n",
-    "\n",
-    "class RNGSubModule(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self):\n",
-    "    print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n",
-    "    print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n",
-    "    RNGSubSubModule()()\n",
-    "\n",
-    "class RNGModule(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self):\n",
-    "    print(f\"RNGModule, count 1: {self.make_rng('rng_stream')}\")\n",
-    "    print(f\"RNGModule, count 2: {self.make_rng('rng_stream')}\")\n",
-    "    RNGSubModule()()\n",
-    "\n",
-    "rng_module = RNGModule()\n",
-    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "As previously discussed, the data that is fed into the Flax hash function consists of:\n",
-    "\n",
-    "  * The path of the `Module`, provided by `self.scope.path` (the top-level root module has an empty path `()`); and\n",
-    "  * The call count for the specific RNG stream.\n",
-    "\n",
-    "In addition, note that each Flax `Module` and sub-`Module` have their own individual call counts, even for the same RNG stream. The convention for sub-`Module` names is: `f'{module_name}_{module_number}'`. For example, the first `Dense` sub-`Module` will be called `Dense_0`, the second one will be called `Dense_1`, and so on.\n",
-    "\n",
-    "Therefore, the following data will be fed into the hash function:\n",
-    "\n",
-    "  * For `RNGModule`: The data is just the call count, such as `(1,)` and `(2,)`, since the root `Module` has an empty path.\n",
-    "  * For `RNGSubModule`: The data is `('RNGSubModule_0', 1)` and `('RNGSubModule_0', 2)`.\n",
-    "  * For `RNGSubSubModule`: The data is `('RNGSubModule_0', 'RNGSubSubModule_0', 1)` and `('RNGSubModule_0', 'RNGSubSubModule_0', 2)`.\n",
-    "\n",
-    "With this data, you can manually reproduce the PRNG keys generated from the `Module` and sub-`Module`s using `self.make_rng`.\n",
-    "\n",
-    "For example:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 55,
-   "metadata": {
-    "outputId": "d38896e6-8061-4f54-fa30-680b3f524071"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1428664606 3351135085]\n",
-      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[3456700291 3873160899]\n",
-      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[3858825717 2323087578]\n",
-      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[ 601859108 3782857444]\n",
-      "RNGSubSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[ 234240654 1028548813]\n",
-      "RNGSubSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[3650462303 2124609379]\n"
-     ]
-    }
-   ],
-   "source": [
-    "stream_seed = jax.random.key(0)\n",
-    "for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_0', 'RNGSubSubModule_0')):\n",
-    "  if initial_data:\n",
-    "    module_name = initial_data[-1]\n",
-    "  else:\n",
-    "    module_name = 'RNGModule'\n",
-    "  for call_count in (1, 2):\n",
-    "    hash_int = produce_hash(data=initial_data+(call_count,))\n",
-    "    rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\n",
-    "    print(f\"{module_name}, count {call_count}: {rng_key}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "If the same sub-`Module` class is used multiple times, you can increment the suffix of the sub-`Module` name accordingly. For example: `RNGSubModule_0`, `RNGSubModule_1`, and so on."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 56,
-   "metadata": {
-    "outputId": "556a79c6-0b38-4736-f5ad-79fd1976b191"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1543086838 3704909070]\n",
-      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[2702764981 3978623664]\n",
-      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[3619592043  626287670]\n",
-      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[965377860 480622172]\n",
-      "RNGSubModule_1, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[2339259155 3791202376]\n",
-      "RNGSubModule_1, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[3825254750  938582370]\n"
-     ]
-    }
-   ],
-   "source": [
-    "class RNGSubModule(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self):\n",
-    "    print(f\"{self.name}, count 1: {self.make_rng('rng_stream')}\")\n",
-    "    print(f\"{self.name}, count 2: {self.make_rng('rng_stream')}\")\n",
-    "\n",
-    "class RNGModule(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self):\n",
-    "    print(f\"RNGModule, count 1: {self.make_rng('rng_stream')}\")\n",
-    "    print(f\"RNGModule, count 2: {self.make_rng('rng_stream')}\")\n",
-    "    RNGSubModule()()\n",
-    "    RNGSubModule()()\n",
-    "\n",
-    "rng_module = RNGModule()\n",
-    "variables = rng_module.init({'rng_stream': jax.random.key(0)})"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 57,
-   "metadata": {
-    "outputId": "433f5e48-d379-490f-e499-ef4c24032776"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "RNGModule, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[1428664606 3351135085]\n",
-      "RNGModule, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[3456700291 3873160899]\n",
-      "RNGSubModule_0, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[3858825717 2323087578]\n",
-      "RNGSubModule_0, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[ 601859108 3782857444]\n",
-      "RNGSubModule_1, count 1: Array((), dtype=key<fry>) overlaying:\n",
-      "[ 426957352 2006350344]\n",
-      "RNGSubModule_1, count 2: Array((), dtype=key<fry>) overlaying:\n",
-      "[4006253729 4205356731]\n"
-     ]
-    }
-   ],
-   "source": [
-    "stream_seed = jax.random.key(0)\n",
-    "for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_1',)):\n",
-    "  if initial_data:\n",
-    "    module_name = initial_data[-1]\n",
-    "  else:\n",
-    "    module_name = 'RNGModule'\n",
-    "  for call_count in (1, 2):\n",
-    "    hash_int = produce_hash(data=initial_data+(call_count,))\n",
-    "    rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))\n",
-    "    print(f\"{module_name}, count {call_count}: {rng_key}\")"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Using `self.param` and `self.variable`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Flax users have the option of creating additional parameters and variables in their modules by using the [`self.param`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.param) and [`self.variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.variable) `Module` methods. An `init_fn` argument must be passed to these methods so that it can generate the initial value of the parameter/variable. `self.make_rng` is commonly used implicitly or explicitly in this `init_fn`, since many initializer functions are stochastic in nature and require a PRNG key. See the full list of Flax initializers [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/initializers.html).\n",
-    "\n",
-    "There are a couple of differences between the two methods that the user should take note of:\n",
-    "* `self.param` always creates a parameter in the `'params'` [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Params-parameters), whereas `self.variable` creates a variable in any [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Variable-collections) the user specifies\n",
-    "* `self.param` will automatically call `self.make_rng('params')` and pass in the generated PRNG key implicitly to the `init_fn` of the parameter you instantiated (it will be passed in as the first argument), whereas users will have to manually specify what RNG stream to call `self.make_rng` on in the `init_fn` of `self.variable` (it could be `'params'` or something different).\n",
-    "\n",
-    "Below is an example using both `self.param` and `self.variable`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 58,
-   "metadata": {
-    "outputId": "fc7df4d6-5d6e-4c3c-98f2-d5b0e9a8a45e"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[[ 0.6227115  -0.0870162 ]\n",
-      " [-0.35369048  0.68549377]]\n",
-      "[[-1.7602162  -0.807241  ]\n",
-      " [-0.09234211  0.79185605]]\n",
-      "[[-1.5997566  -0.09624068]\n",
-      " [ 0.21725878  0.16356604]]\n",
-      "[[ 0.44956833 -1.1854612 ]\n",
-      " [ 0.17371362 -0.768862  ]]\n",
-      "[[ 1.5019631  -1.3983697 ]\n",
-      " [-0.48637655 -1.4504721 ]]\n",
-      "[ 2.114508  -0.5090715]\n"
-     ]
-    }
-   ],
-   "source": [
-    "class Model(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x):\n",
-    "    # kernel will use 'params' seed, initial data will include 'Dense_0', call count 1\n",
-    "    x = nn.Dense(2, kernel_init=jax.random.normal, use_bias=False)(x)\n",
-    "    # model_param will use 'params' seed, call count 1\n",
-    "    model_param = self.param('model_param', jax.random.normal, x.shape)\n",
-    "    # model_variable1 will use 'params' seed, call count 2\n",
-    "    model_variable1 = self.variable(\n",
-    "      'other_collection',\n",
-    "      'model_variable1',\n",
-    "      lambda: jax.random.normal(self.make_rng('params'), x.shape),\n",
-    "    )\n",
-    "    # model_variable2 will use 'other' seed, call count 1\n",
-    "    model_variable2 = self.variable(\n",
-    "      'other_collection',\n",
-    "      'model_variable2',\n",
-    "      lambda: jax.random.normal(self.make_rng('other'), x.shape),\n",
-    "    )\n",
-    "    # kernel will use 'params' seed, initial data will include 'Dense_1', call count 1\n",
-    "    # bias will use 'params' seed, initial data will include 'Dense_1', call count 2\n",
-    "    x = nn.Dense(2, kernel_init=jax.random.normal, bias_init=jax.random.normal)(\n",
-    "      x\n",
-    "    )\n",
-    "    return x\n",
-    "\n",
-    "model = Model()\n",
-    "variables = model.init(\n",
-    "  {'params': jax.random.key(0), 'other': jax.random.key(1)}, jnp.ones((2, 2))\n",
-    ")\n",
-    "print(variables['params']['Dense_0']['kernel'])\n",
-    "print(variables['params']['model_param'])\n",
-    "print(variables['other_collection']['model_variable1'])\n",
-    "print(variables['other_collection']['model_variable2'])\n",
-    "print(variables['params']['Dense_1']['kernel'])\n",
-    "print(variables['params']['Dense_1']['bias'])"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Remember:\n",
-    "* there is a separate count for each RNG stream; this is why the count for `self.make_rng('other')` starts at 1 even though there were earlier calls of `self.make_rng('params')`\n",
-    "* each submodule has their own separate count for each rng stream; this is why each `Dense` layer has their own separate count for `self.make_rng('params')` and why `model_param` and `model_variable1` share the same count (since they are defined within the same top-level parent module)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 59,
-   "metadata": {
-    "outputId": "40745795-8bbc-4c31-811a-f3658e5459d7"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[[-1.6185919   0.700908  ]\n",
-      " [-1.3146383  -0.79342234]]\n",
-      "[[ 0.0761425 -1.6157459]\n",
-      " [-1.6857724  0.7126891]]\n",
-      "[[ 0.60175574  0.2553228 ]\n",
-      " [ 0.27367848 -2.1975214 ]]\n",
-      "[[1.6249592  0.30813068]\n",
-      " [1.6613585  1.0404155 ]]\n",
-      "[[ 0.0030665   0.29551846]\n",
-      " [ 0.16670242 -0.78252524]]\n",
-      "[[1.582462   0.15216611]]\n"
-     ]
-    }
-   ],
-   "source": [
-    "params_seed = jax.random.key(0)\n",
-    "other_seed = jax.random.key(1)\n",
-    "for initial_data, count, seed, shape in (\n",
-    "  (('Dense_0',), 1, params_seed, (2, 2)),\n",
-    "  ((), 1, params_seed, (2, 2)),\n",
-    "  ((), 2, params_seed, (2, 2)),\n",
-    "  ((), 1, other_seed, (2, 2)),\n",
-    "  (('Dense_1',), 1, params_seed, (2, 2)),\n",
-    "  (('Dense_1',), 2, params_seed, (1, 2)),\n",
-    "):\n",
-    "  hash_int = produce_hash(data=(*initial_data, count))\n",
-    "  rng_key = jax.random.fold_in(seed, jnp.uint32(hash_int))\n",
-    "  print(jax.random.normal(rng_key, shape))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Managing RNG streams inside a training loop"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Below is an example of managing RNG streams from `self.make_rng`, `self.param`, `self.variable` and `nn.Dropout` in a training loop (note: `nn.Dropout` requires a seed PRNG key to be passed in the `'dropout'` RNG stream, since it implicitly calls `self.make_rng('dropout')`):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 60,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "class SubModule(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x, train):\n",
-    "    # variables created using `self.param` will use `self.make_rng('params')`\n",
-    "    kernel = self.param('submodule_kernel', jax.random.normal, x.shape)\n",
-    "    x = x + kernel\n",
-    "    # `nn.Dropout` will use self.make_rng('dropout')\n",
-    "    x = nn.Dropout(0.2)(x, deterministic=not train)\n",
-    "    # `nn.Dense` will use self.make_rng('params')\n",
-    "    x = nn.Dense(3)(x)\n",
-    "    return x\n",
-    "\n",
-    "class Model(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x, train):\n",
-    "    # make kernel use `self.make_rng('other')`\n",
-    "    kernel = self.variable(\n",
-    "      'other_collection',\n",
-    "      'module_kernel',\n",
-    "      lambda: jax.random.normal(self.make_rng('other'), x.shape),\n",
-    "    )\n",
-    "    x = (\n",
-    "      x + kernel.value\n",
-    "    )  # `.value` will extract the underlying value of the variable\n",
-    "    x = SubModule()(x, train)\n",
-    "    # `nn.Dropout` will use self.make_rng('dropout')\n",
-    "    x = nn.Dropout(0.2)(x, deterministic=not train)\n",
-    "    # `nn.Dense` will use self.make_rng('params')\n",
-    "    x = nn.Dense(2)(x)\n",
-    "    return x\n",
-    "\n",
-    "params_rng, other_rng, train_rng = jax.random.split(jax.random.key(0), 3)\n",
-    "init_rngs = {'params': params_rng, 'other': other_rng}\n",
-    "\n",
-    "x = jnp.ones((1, 3))\n",
-    "y = jnp.ones((1, 2))\n",
-    "\n",
-    "module = Model()\n",
-    "variables = module.init(init_rngs, x, train=False)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 61,
-   "metadata": {
-    "outputId": "9e689bfd-ba73-4b78-c3c0-89693caada2d"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "0.86758584\n",
-      "0.86374676\n",
-      "0.861765\n",
-      "0.85973275\n",
-      "0.8573686\n",
-      "0.85558856\n",
-      "0.8524709\n",
-      "0.84871066\n",
-      "0.8456431\n",
-      "0.84329045\n"
-     ]
-    }
-   ],
-   "source": [
-    "def update(variables, rng):\n",
-    "  # we don't need to provide a 'params' or 'other' rng, as only 'dropout' rng will be used during training\n",
-    "  # split the rng to get a dropout_rng to be used for this training iteration,\n",
-    "  # and to get another rng key to be used for the next training iteration\n",
-    "  dropout_rng, next_rng = jax.random.split(rng)\n",
-    "  def loss(params):\n",
-    "    out = module.apply(\n",
-    "      {'params': params, 'other_collection': variables['other_collection']},\n",
-    "      x,\n",
-    "      train=True,\n",
-    "      rngs={'dropout': dropout_rng},\n",
-    "    )\n",
-    "    return jnp.mean((y - out) ** 2)\n",
-    "  grads = jax.grad(loss)(variables['params'])\n",
-    "  params = jax.tree_map(lambda p, g: p - 1e-3 * g, variables['params'], grads)\n",
-    "  return {\n",
-    "    'params': params,\n",
-    "    'other_collection': variables['other_collection'],\n",
-    "  }, next_rng\n",
-    "\n",
-    "for _ in range(10):\n",
-    "  variables, train_rng = update(variables, train_rng)\n",
-    "  out = module.apply(variables, x, train=False)\n",
-    "  print(jnp.mean((y - out)**2))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### 🔪 Sharp edge 🔪 - unintentionally generating the same values\n",
-    "\n",
-    "There is an edge case where the same value can be unintentionally generated.\n",
-    "See the [Flax issue](https://github.com/google/flax/issues/2157) for more details."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 62,
-   "metadata": {
-    "outputId": "80e45c76-bb4d-48ab-8fac-8b14126428d1"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array(False, dtype=bool)"
-      ]
-     },
-     "execution_count": 62,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "class Leaf(nn.Module):\n",
-    "  def __call__(self, x):\n",
-    "    return x + jax.random.randint(self.make_rng(\"rng\"), (), 0, 100)\n",
-    "\n",
-    "class Node(nn.Module):\n",
-    "  leaf_name: str\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x):\n",
-    "    return Leaf(name=self.leaf_name)(x)\n",
-    "\n",
-    "class Model(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x):\n",
-    "    return (Node(name=\"ab\", leaf_name=\"cdef\")(x),\n",
-    "            Node(name=\"abc\", leaf_name=\"def\")(x),\n",
-    "    )\n",
-    "\n",
-    "out1, out2 = Model().apply({}, 0, rngs={\"rng\": jax.random.key(33)})\n",
-    "out1 == out2 # same output, despite having different submodule names"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This occurs because the hash function [concatenates strings together](https://docs.python.org/3/library/hashlib.html#hashlib.hash.update), so the data `('AB', 'C')` is equivalent to data `('A', 'BC')` when fed into the hash function, therefore producing the same hash int."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 63,
-   "metadata": {
-    "outputId": "07fc8349-41a3-4c74-d25c-0ffc4aa1be0a"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "947574064\n",
-      "947574064\n",
-      "947574064\n",
-      "947574064\n"
-     ]
-    }
-   ],
-   "source": [
-    "print(produce_hash(data=('A', 'B', 'C', 1)))\n",
-    "print(produce_hash(data=('AB', 'C', 1)))\n",
-    "print(produce_hash(data=('A', 'BC', 1)))\n",
-    "print(produce_hash(data=('ABC', 1)))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "To avoid this edge case, users can flip the `flax_fix_rng_separator` [configuration flag](https://flax.readthedocs.io/en/latest/api_reference/flax.config.html#flax.configurations.Config.flax_fix_rng_separator) to `True`."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 64,
-   "metadata": {
-    "outputId": "4ba0c3cb-e903-40af-ab3f-687d83c257c9"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array(False, dtype=bool)"
-      ]
-     },
-     "execution_count": 64,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "flax.config.update('flax_fix_rng_separator', True)\n",
-    "out1, out2 = Model().apply({}, 0, rngs={\"rng\": jax.random.key(33)})\n",
-    "out1 == out2 # different output"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Managing RNG's on multiple devices"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "This section will show examples on how to use `jit` and `shard_map` to use RNG's in multi-device settings."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Using `jax.jit`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "When using [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), we can use RNG's as we did before, but we now include `in_shardings` and `out_shardings` arguments to shard specify how to shard input and output data.\n",
-    "\n",
-    "For more details on training on multiple devices in Flax using `jax.jit`, see our [Scale up Flax Modules on multiple devices guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html#) and [lm1b example](https://github.com/google/flax/tree/main/examples/lm1b)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 65,
-   "metadata": {
-    "outputId": "9deab9d8-3e15-4be0-8d91-279f6984ee99"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[CpuDevice(id=0) CpuDevice(id=1) CpuDevice(id=2) CpuDevice(id=3)\n",
-      " CpuDevice(id=4) CpuDevice(id=5) CpuDevice(id=6) CpuDevice(id=7)]\n",
-      "Mesh(device_ids=array([0, 1, 2, 3, 4, 5, 6, 7]), axis_names=('data',))\n",
-      "NamedSharding(mesh=Mesh('data': 8), spec=PartitionSpec('data',))\n"
-     ]
-    }
-   ],
-   "source": [
-    "# Create a mesh and annotate the axis with a name.\n",
-    "device_mesh = mesh_utils.create_device_mesh((8,))\n",
-    "print(device_mesh)\n",
-    "\n",
-    "mesh = Mesh(devices=device_mesh, axis_names=('data',))\n",
-    "print(mesh)\n",
-    "\n",
-    "data_sharding = NamedSharding(mesh, PartitionSpec('data',))\n",
-    "print(data_sharding)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 66,
-   "metadata": {
-    "outputId": "9b47eef2-612e-4d9f-ffb2-1e868cb52d86"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([[-2.2187614 ],\n",
-       "       [-2.8055234 ],\n",
-       "       [-2.5464187 ],\n",
-       "       [ 1.0270392 ],\n",
-       "       [-3.5243359 ],\n",
-       "       [-2.2795477 ],\n",
-       "       [-0.6504516 ],\n",
-       "       [ 0.17373264]], dtype=float32)"
-      ]
-     },
-     "execution_count": 66,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "class Model(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x, add_noise):\n",
-    "    x = nn.Dense(1)(x)\n",
-    "    # use jnp.where for control flow; for more details see: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError\n",
-    "    return jnp.where(\n",
-    "      add_noise, x + jax.random.normal(self.make_rng('params'), x.shape), x\n",
-    "    )\n",
-    "\n",
-    "module = Model()\n",
-    "init_rng, apply_rng = jax.random.split(jax.random.key(0))\n",
-    "x = jnp.ones((8, 1))\n",
-    "variables = module.init(init_rng, x, False)\n",
-    "\n",
-    "# create custom forward function, since jit does not support kwargs when in_shardings is specified\n",
-    "def forward(variables, x, add_noise, rng):\n",
-    "  return module.apply(variables, x, add_noise, rngs={'params': rng})\n",
-    "\n",
-    "# shard the inputs x across devices\n",
-    "# replicate the variables, add_noise boolean and rng key across devices\n",
-    "# shard the output across devices\n",
-    "jit_forward = jax.jit(\n",
-    "  forward,\n",
-    "  in_shardings=(None, data_sharding, None, None),\n",
-    "  out_shardings=data_sharding,\n",
-    ")\n",
-    "out = jit_forward(variables, x, True, apply_rng)\n",
-    "out"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "The output is different given the same input, meaning the RNG key was used to add noise to the output.\n",
-    "\n",
-    "We can also confirm that the output is sharded across devices:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 67,
-   "metadata": {
-    "outputId": "f6596bdc-89b4-46bf-ba99-84b1009d8156"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[Shard(device=CpuDevice(id=0), index=(slice(0, 1, None), slice(None, None, None)), replica_id=0, data=[[-2.2187614]]),\n",
-       " Shard(device=CpuDevice(id=1), index=(slice(1, 2, None), slice(None, None, None)), replica_id=0, data=[[-2.8055234]]),\n",
-       " Shard(device=CpuDevice(id=2), index=(slice(2, 3, None), slice(None, None, None)), replica_id=0, data=[[-2.5464187]]),\n",
-       " Shard(device=CpuDevice(id=3), index=(slice(3, 4, None), slice(None, None, None)), replica_id=0, data=[[1.0270392]]),\n",
-       " Shard(device=CpuDevice(id=4), index=(slice(4, 5, None), slice(None, None, None)), replica_id=0, data=[[-3.5243359]]),\n",
-       " Shard(device=CpuDevice(id=5), index=(slice(5, 6, None), slice(None, None, None)), replica_id=0, data=[[-2.2795477]]),\n",
-       " Shard(device=CpuDevice(id=6), index=(slice(6, 7, None), slice(None, None, None)), replica_id=0, data=[[-0.6504516]]),\n",
-       " Shard(device=CpuDevice(id=7), index=(slice(7, 8, None), slice(None, None, None)), replica_id=0, data=[[0.17373264]])]"
-      ]
-     },
-     "execution_count": 67,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out.addressable_shards"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Another way to visualize the output sharding:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 68,
-   "metadata": {
-    "outputId": "7df87ccd-0979-4b85-fbac-1d31aac53276"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "┌───────┐\n",
-      "│ CPU 0 │\n",
-      "├───────┤\n",
-      "│ CPU 1 │\n",
-      "├───────┤\n",
-      "│ CPU 2 │\n",
-      "├───────┤\n",
-      "│ CPU 3 │\n",
-      "├───────┤\n",
-      "│ CPU 4 │\n",
-      "├───────┤\n",
-      "│ CPU 5 │\n",
-      "├───────┤\n",
-      "│ CPU 6 │\n",
-      "├───────┤\n",
-      "│ CPU 7 │\n",
-      "└───────┘\n"
-     ]
-    }
-   ],
-   "source": [
-    "jax.debug.visualize_array_sharding(out)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "If we choose not to add noise, then the output is the same across all batches (as expected, since the input is the same for all batches):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 69,
-   "metadata": {
-    "outputId": "cb7eb7d0-3fa7-4f72-a129-18901dc61bb1"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "[[-1.2839764]\n",
-      " [-1.2839764]\n",
-      " [-1.2839764]\n",
-      " [-1.2839764]\n",
-      " [-1.2839764]\n",
-      " [-1.2839764]\n",
-      " [-1.2839764]\n",
-      " [-1.2839764]]\n"
-     ]
-    }
-   ],
-   "source": [
-    "out = jit_forward(variables, x, False, apply_rng)\n",
-    "print(out)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can confirm the un-jitted function produces the same values, albeit unsharded (note there may be small numerical differences due to compiler optimizations from jitting):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 70,
-   "metadata": {
-    "outputId": "b88834c1-ebb2-422a-a665-1732015a3974"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([[-2.2187614 ],\n",
-       "       [-2.8055234 ],\n",
-       "       [-2.5464187 ],\n",
-       "       [ 1.0270392 ],\n",
-       "       [-3.5243359 ],\n",
-       "       [-2.2795477 ],\n",
-       "       [-0.6504516 ],\n",
-       "       [ 0.17373264]], dtype=float32)"
-      ]
-     },
-     "execution_count": 70,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out = forward(variables, x, True, apply_rng)\n",
-    "out"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 71,
-   "metadata": {
-    "outputId": "5b06f58b-7f1f-4cd6-c29f-3da5c0ea189f"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([[-1.2839764],\n",
-       "       [-1.2839764],\n",
-       "       [-1.2839764],\n",
-       "       [-1.2839764],\n",
-       "       [-1.2839764],\n",
-       "       [-1.2839764],\n",
-       "       [-1.2839764],\n",
-       "       [-1.2839764]], dtype=float32)"
-      ]
-     },
-     "execution_count": 71,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out = forward(variables, x, False, apply_rng)\n",
-    "out"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### Using `shard_map`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "When using [`jax.experimental.shard_map.shard_map`](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html), the important parts to remember are to:\n",
-    "* split your PRNG key to produce a different key for each device\n",
-    "* the PRNG keys will be sharded automatically to each device (provided you use the correct partition specification), but the [**rank of the original batched PRNG key array will not be reduced**](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html#rank-reducing-vs-rank-preserving-maps-over-array-axes); e.g.\n",
-    "with a batch of 8 PRNG keys and 8 devices, each device will see a PRNG key batch of size 1 within the `shard_map`-ed function\n",
-    "  * therefore to access the PRNG key itself, we need to index slice into it (see the example below)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 72,
-   "metadata": {
-    "outputId": "29b2911f-28a7-4a25-fcf7-4c3544dafada"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "Array([[-1.2605132 ],\n",
-       "       [-1.2405176 ],\n",
-       "       [-0.99350417],\n",
-       "       [-1.0277128 ],\n",
-       "       [-1.4154483 ],\n",
-       "       [-0.3905797 ],\n",
-       "       [-2.417677  ],\n",
-       "       [ 0.9023453 ]], dtype=float32)"
-      ]
-     },
-     "execution_count": 72,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "def forward(variables, x, add_noise, rng_key_batch):\n",
-    "  # rng_key_batch is a batch of size 1 containing 1 PRNG key\n",
-    "  # index slice into the rng_key_batch to access the PRNG key\n",
-    "  return module.apply(\n",
-    "    variables, x, add_noise, rngs={'params': rng_key_batch[0]}\n",
-    "  )\n",
-    "\n",
-    "# define partition specifications\n",
-    "data_pspec = PartitionSpec('data')\n",
-    "no_pspec = PartitionSpec()\n",
-    "\n",
-    "# shard the inputs x and rng keys across devices\n",
-    "# replicate the variables and add_noise boolean across devices\n",
-    "# shard the output across devices\n",
-    "shmap_forward = shard_map(\n",
-    "  forward,\n",
-    "  mesh=mesh,\n",
-    "  in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),\n",
-    "  out_specs=data_pspec,\n",
-    ")\n",
-    "# get 8 different rng's that will be used by the 8 devices when doing forward inference\n",
-    "apply_rngs = jax.random.split(apply_rng, 8)\n",
-    "out = shmap_forward(variables, x, True, apply_rngs)\n",
-    "out"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Confirm that the output is sharded across devices:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 73,
-   "metadata": {
-    "outputId": "2d74dff9-3a7d-47df-afbc-3bda4b8edf51"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "[Shard(device=CpuDevice(id=0), index=(slice(0, 1, None), slice(None, None, None)), replica_id=0, data=[[-1.2605132]]),\n",
-       " Shard(device=CpuDevice(id=1), index=(slice(1, 2, None), slice(None, None, None)), replica_id=0, data=[[-1.2405176]]),\n",
-       " Shard(device=CpuDevice(id=2), index=(slice(2, 3, None), slice(None, None, None)), replica_id=0, data=[[-0.99350417]]),\n",
-       " Shard(device=CpuDevice(id=3), index=(slice(3, 4, None), slice(None, None, None)), replica_id=0, data=[[-1.0277128]]),\n",
-       " Shard(device=CpuDevice(id=4), index=(slice(4, 5, None), slice(None, None, None)), replica_id=0, data=[[-1.4154483]]),\n",
-       " Shard(device=CpuDevice(id=5), index=(slice(5, 6, None), slice(None, None, None)), replica_id=0, data=[[-0.3905797]]),\n",
-       " Shard(device=CpuDevice(id=6), index=(slice(6, 7, None), slice(None, None, None)), replica_id=0, data=[[-2.417677]]),\n",
-       " Shard(device=CpuDevice(id=7), index=(slice(7, 8, None), slice(None, None, None)), replica_id=0, data=[[0.9023453]])]"
-      ]
-     },
-     "execution_count": 73,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "out.addressable_shards"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 74,
-   "metadata": {
-    "outputId": "28bfe59b-ded1-455a-c894-f1c457ec28bf"
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "┌───────┐\n",
-      "│ CPU 0 │\n",
-      "├───────┤\n",
-      "│ CPU 1 │\n",
-      "├───────┤\n",
-      "│ CPU 2 │\n",
-      "├───────┤\n",
-      "│ CPU 3 │\n",
-      "├───────┤\n",
-      "│ CPU 4 │\n",
-      "├───────┤\n",
-      "│ CPU 5 │\n",
-      "├───────┤\n",
-      "│ CPU 6 │\n",
-      "├───────┤\n",
-      "│ CPU 7 │\n",
-      "└───────┘\n"
-     ]
-    }
-   ],
-   "source": [
-    "jax.debug.visualize_array_sharding(out)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Lifted transforms"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "[Flax lifted transforms](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html) allow you to use [JAX transforms](https://github.com/google/jax#transformations) with `Module` arguments. This section will show you how to control how PRNG keys are split in Flax lifted transforms.\n",
-    "\n",
-    "Refer to [Lifted transformations](https://flax.readthedocs.io/en/latest/developer_notes/lift.html) for more detail."
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### `nn.vmap`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can use [`nn.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.vmap.html) to create a batched `Dense` layer:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 75,
-   "metadata": {
-    "outputId": "7d7c4c3c-2ed2-40dd-f8d3-bb1ef91d93eb"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'params': {'bias': Array([0., 0.], dtype=float32),\n",
-       "  'kernel': Array([[-1.2488099 , -0.6127134 ],\n",
-       "         [-0.07084481,  0.60130936]], dtype=float32)}}"
-      ]
-     },
-     "execution_count": 75,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "x = jnp.ones((3, 2))\n",
-    "\n",
-    "BatchDense = nn.vmap(\n",
-    "    nn.Dense,\n",
-    "    in_axes=0, out_axes=0,\n",
-    "    variable_axes={'params': None},\n",
-    "    split_rngs={'params': False})\n",
-    "\n",
-    "BatchDense(2).init(jax.random.key(0), x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "By denoting `variable_axes={'params': 0}'`, we vectorize the `params` Arrays on the first axis. However the parameter values generated are all identical to each other:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 76,
-   "metadata": {
-    "outputId": "93152c8f-ada9-4cf5-c0e2-560a284b3981"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'params': {'bias': Array([[0., 0.],\n",
-       "         [0., 0.],\n",
-       "         [0., 0.]], dtype=float32),\n",
-       "  'kernel': Array([[[-1.2488099 , -0.6127134 ],\n",
-       "          [-0.07084481,  0.60130936]],\n",
-       "  \n",
-       "         [[-1.2488099 , -0.6127134 ],\n",
-       "          [-0.07084481,  0.60130936]],\n",
-       "  \n",
-       "         [[-1.2488099 , -0.6127134 ],\n",
-       "          [-0.07084481,  0.60130936]]], dtype=float32)}}"
-      ]
-     },
-     "execution_count": 76,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "BatchDense = nn.vmap(\n",
-    "    nn.Dense,\n",
-    "    in_axes=0, out_axes=0,\n",
-    "    variable_axes={'params': 0},\n",
-    "    split_rngs={'params': False})\n",
-    "\n",
-    "BatchDense(2).init(jax.random.key(0), x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "If we also make `split_rngs={'params': True}`, then the PRNG key we provide is split across the variable axis (in this case, the batch axis 0), and we can generate different parameters for each batch input:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 77,
-   "metadata": {
-    "outputId": "dfbf4e49-55dc-4a6f-bece-396cb48d875e"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'params': {'bias': Array([[0., 0.],\n",
-       "         [0., 0.],\n",
-       "         [0., 0.]], dtype=float32),\n",
-       "  'kernel': Array([[[-0.2526208 , -0.15088455],\n",
-       "          [-1.1987205 , -0.40843305]],\n",
-       "  \n",
-       "         [[-0.7064888 , -1.108805  ],\n",
-       "          [-0.938775  ,  1.4812315 ]],\n",
-       "  \n",
-       "         [[-0.59468937, -0.2502723 ],\n",
-       "          [-1.33515   ,  0.5067442 ]]], dtype=float32)}}"
-      ]
-     },
-     "execution_count": 77,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "BatchDense = nn.vmap(\n",
-    "    nn.Dense,\n",
-    "    in_axes=0, out_axes=0,\n",
-    "    variable_axes={'params': 0},\n",
-    "    split_rngs={'params': True})\n",
-    "\n",
-    "BatchDense(2).init(jax.random.key(0), x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Adding a variable via `self.variable` is straightforward:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 78,
-   "metadata": {
-    "outputId": "20fa39a3-9e7a-4a58-8ad7-48a0119ec466"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'other_collection': {'kernel': Array([[-0.8193048 ,  0.711106  ],\n",
-       "         [-0.37802765, -0.66705877],\n",
-       "         [-0.44808003,  0.93031347]], dtype=float32)},\n",
-       " 'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
-       "          [0., 0.],\n",
-       "          [0., 0.]], dtype=float32),\n",
-       "   'kernel': Array([[[-0.9079084 ,  0.76390624],\n",
-       "           [-0.01285526,  0.4320353 ]],\n",
-       "   \n",
-       "          [[ 0.12398645,  0.7884565 ],\n",
-       "           [ 1.5344163 ,  1.3186085 ]],\n",
-       "   \n",
-       "          [[-0.44171348,  0.43430036],\n",
-       "           [-0.40732604,  0.29774475]]], dtype=float32)}}}"
-      ]
-     },
-     "execution_count": 78,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "class Model(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x):\n",
-    "    x = nn.Dense(2)(x)\n",
-    "    kernel = self.variable(\n",
-    "      'other_collection',\n",
-    "      'kernel',\n",
-    "      lambda: jax.random.normal(self.make_rng('other'), x.shape),\n",
-    "    )\n",
-    "    return x + kernel.value\n",
-    "\n",
-    "BatchModel = nn.vmap(\n",
-    "  Model,\n",
-    "  in_axes=0,\n",
-    "  out_axes=0,\n",
-    "  variable_axes={'params': 0, 'other_collection': 0},\n",
-    "  split_rngs={'params': True, 'other': True},\n",
-    ")\n",
-    "\n",
-    "BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can control which RNG stream to split, for example, if we only wanted to split the `'params'` RNG stream, then the variables generated from `self.variable` will be the same for each batch input:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 79,
-   "metadata": {
-    "outputId": "69d52ec5-9018-43c8-9d77-65e9bf00b79f"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'other_collection': {'kernel': Array([[ 0.44956833, -1.1854612 ],\n",
-       "         [ 0.44956833, -1.1854612 ],\n",
-       "         [ 0.44956833, -1.1854612 ]], dtype=float32)},\n",
-       " 'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
-       "          [0., 0.],\n",
-       "          [0., 0.]], dtype=float32),\n",
-       "   'kernel': Array([[[-0.9079084 ,  0.76390624],\n",
-       "           [-0.01285526,  0.4320353 ]],\n",
-       "   \n",
-       "          [[ 0.12398645,  0.7884565 ],\n",
-       "           [ 1.5344163 ,  1.3186085 ]],\n",
-       "   \n",
-       "          [[-0.44171348,  0.43430036],\n",
-       "           [-0.40732604,  0.29774475]]], dtype=float32)}}}"
-      ]
-     },
-     "execution_count": 79,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "BatchModel = nn.vmap(\n",
-    "    Model,\n",
-    "    in_axes=0, out_axes=0,\n",
-    "    variable_axes={'params': 0, 'other_collection': 0},\n",
-    "    split_rngs={'params': True, 'other': False})\n",
-    "\n",
-    "BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can also control which parameters / variables should be generated for each batch input, for example, if we only wanted `'params'` to generate separate parameters for each batch input:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 80,
-   "metadata": {
-    "outputId": "b01187f9-aded-41bc-d17f-e2965355675a"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'other_collection': {'kernel': Array([ 0.44956833, -1.1854612 ], dtype=float32)},\n",
-       " 'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
-       "          [0., 0.],\n",
-       "          [0., 0.]], dtype=float32),\n",
-       "   'kernel': Array([[[-0.9079084 ,  0.76390624],\n",
-       "           [-0.01285526,  0.4320353 ]],\n",
-       "   \n",
-       "          [[ 0.12398645,  0.7884565 ],\n",
-       "           [ 1.5344163 ,  1.3186085 ]],\n",
-       "   \n",
-       "          [[-0.44171348,  0.43430036],\n",
-       "           [-0.40732604,  0.29774475]]], dtype=float32)}}}"
-      ]
-     },
-     "execution_count": 80,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "BatchModel = nn.vmap(\n",
-    "    Model,\n",
-    "    in_axes=0, out_axes=0,\n",
-    "    variable_axes={'params': 0, 'other_collection': None},\n",
-    "    split_rngs={'params': True, 'other': False})\n",
-    "\n",
-    "BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "### `nn.scan`"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "We can use [`nn.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.scan.html) to create a scanned `Module` layer (this is useful for simplifying repetitively stacked submodules):"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 81,
-   "metadata": {
-    "outputId": "fda568b4-3c1d-4a5e-96f0-058c6fc5b49a"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
-       "          [0., 0.],\n",
-       "          [0., 0.]], dtype=float32),\n",
-       "   'kernel': Array([[[-0.07838312, -0.7422982 ],\n",
-       "           [ 0.87488323,  0.13773395]],\n",
-       "   \n",
-       "          [[ 0.97309333,  0.9087693 ],\n",
-       "           [-0.12564984, -1.0920651 ]],\n",
-       "   \n",
-       "          [[-0.99055105,  1.1499453 ],\n",
-       "           [-0.15721127, -0.62520015]]], dtype=float32)}}}"
-      ]
-     },
-     "execution_count": 81,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "x = jnp.ones((3, 2))\n",
-    "\n",
-    "class ResidualMLPBlock(nn.Module):\n",
-    "  @nn.compact\n",
-    "  def __call__(self, x, _):\n",
-    "    h = nn.Dense(features=2)(x)\n",
-    "    h = nn.relu(h)\n",
-    "    return x + h, None # return an empty carry\n",
-    "\n",
-    "ScanMLP = nn.scan(\n",
-    "      ResidualMLPBlock, variable_axes={'params': 0},\n",
-    "      variable_broadcast=False, split_rngs={'params': True},\n",
-    "      length=3)\n",
-    "\n",
-    "ScanMLP().init(jax.random.key(0), x, None) # pass in an empty carry"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Similar to before, we can control whether to split the RNG stream or not, for example, if we wanted all the stacked modules to be initialized to the same parameter values, we can pass in `split_rngs={'params': False}`:"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 82,
-   "metadata": {
-    "outputId": "620aff62-e20d-4794-b3af-5a5058c2471d"
-   },
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "{'params': {'Dense_0': {'bias': Array([[0., 0.],\n",
-       "          [0., 0.],\n",
-       "          [0., 0.]], dtype=float32),\n",
-       "   'kernel': Array([[[-0.66715515, -0.0484313 ],\n",
-       "           [ 0.9867164 ,  0.75408363]],\n",
-       "   \n",
-       "          [[-0.66715515, -0.0484313 ],\n",
-       "           [ 0.9867164 ,  0.75408363]],\n",
-       "   \n",
-       "          [[-0.66715515, -0.0484313 ],\n",
-       "           [ 0.9867164 ,  0.75408363]]], dtype=float32)}}}"
-      ]
-     },
-     "execution_count": 82,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "ScanMLP = nn.scan(\n",
-    "      ResidualMLPBlock, variable_axes={'params': 0},\n",
-    "      variable_broadcast=False, split_rngs={'params': False},\n",
-    "      length=3)\n",
-    "\n",
-    "ScanMLP().init(jax.random.key(0), x, None)"
-   ]
-  }
- ],
- "metadata": {
-  "jupytext": {
-   "formats": "ipynb,md:myst"
-  },
-  "language_info": {
-   "name": "python"
-  }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/docs/guides/flax_fundamentals/rng_guide.md b/docs/guides/flax_fundamentals/rng_guide.md
deleted file mode 100644
index 5bdba844..00000000
--- a/docs/guides/flax_fundamentals/rng_guide.md
+++ /dev/null
@@ -1,797 +0,0 @@
----
-jupytext:
-  formats: ipynb,md:myst
-  text_representation:
-    extension: .md
-    format_name: myst
-    format_version: 0.13
-    jupytext_version: 1.13.8
----
-
-# Randomness and PRNGs in Flax
-
-+++
-
-In this guide, you will learn how Flax uses [JAX's explicit pseudorandom number generator (PRNG) keys](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#jax-prng) to emulate randomness, and adds some additional features to make it easier for users to thread PRNG keys through different Flax `Module`s.
-
-If you are new to JAX PRNG keys or need a refresher, check out:
-- [JAX 101: PRNGs in JAX](https://jax.readthedocs.io/en/latest/jax-101/05-random-numbers.html)
-- [🔪 JAX - The Sharp Bits 🔪: Random Numbers](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#random-numbers).
-
-+++
-
-## Setup
-
-+++
-
-Install or upgrade Flax, and then import some necessary dependencies.
-
-**Note:** This guide uses the `--xla_force_host_platform_device_count=8` flag to emulate multiple devices in a CPU environment in a Google Colab/Jupyter Notebook. You don’t need this if you are already using a multi-device Google Cloud TPU environment, for example, on Google Cloud or in a Kaggle VM with a TPU.
-
-```{code-cell}
-:outputId: bb13d0ba-f908-4746-e4d3-f5916e0670ff
-:tags: [skip-execution]
-
-!pip install -q flax
-```
-
-```{code-cell}
-import os
-os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'
-```
-
-```{code-cell}
-import flax, flax.linen as nn
-import jax, jax.numpy as jnp
-from jax.sharding import Mesh, PartitionSpec, NamedSharding
-from jax.experimental import mesh_utils
-from jax.experimental.shard_map import shard_map
-
-import hashlib
-```
-
-```{code-cell}
-:outputId: ab5e0d3b-2d51-44ee-b823-d152ad1a10b2
-
-jax.devices()
-```
-
-Set the JAX config variable `jax_threefry_partitionable` to `True`. This will be the default value in the future and makes the PRNG more efficiently auto-parallelizable under `jax.jit`. Refer to [JAX discussion](https://github.com/google/jax/discussions/18480) for more details.
-
-```{code-cell}
-jax.config.update('jax_threefry_partitionable', True)
-assert jax.config.jax_threefry_partitionable == True
-assert jax.config.jax_default_prng_impl == 'threefry2x32'
-```
-
-## Receiving, manipulating and creating PRNG keys with `Module.make_rng`
-
-+++
-
-The primary method Flax uses to receive, manipulate and create PRNG keys is via the `Module` method [`self.make_rng`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.make_rng). It is a method that accepts a string name that represents an "RNG stream". Each RNG stream has an initial starting seed PRNG key, which the user passes in as a dictionary argument (i.e. into an [`.init`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.init) or [`.apply`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.apply) function), and the starting seed is used by `self.make_rng` to generate more PRNG keys for that stream.
-
-Note that this method can only be called with bounded modules (see [The Flax Module lifecycle](https://flax.readthedocs.io/en/latest/developer_notes/module_lifecycle.html#top-level-modules)).
-
-```{code-cell}
-:outputId: 8304fbc2-ab0a-42f0-9d4c-c88a74012e83
-
-class RNGModule(nn.Module):
-  @nn.compact
-  def __call__(self):
-    print(self.make_rng('rng_stream'))
-    print(self.make_rng('rng_stream'))
-    print(self.make_rng('rng_stream'))
-
-rng_module = RNGModule()
-variables = rng_module.init({'rng_stream': jax.random.key(0)})
-```
-
-Now if we use a different starting seed PRNG key, we will generate different values (as intended).
-
-```{code-cell}
-:outputId: 1cfd9632-11cf-43f4-b922-63038705a195
-
-variables = rng_module.init({'rng_stream': jax.random.key(1)})
-```
-
-Calling `self.make_rng` for one stream will not affect the random values generated from another stream; i.e. the call order doesn't matter.
-
-```{code-cell}
-:outputId: 60dde8ff-08a9-4743-c7ec-c997aa0537b6
-
-class RNGModuleTwoStreams(nn.Module):
-  @nn.compact
-  def __call__(self):
-    # same value as first code snippet above
-    print(f"rng_stream1: {self.make_rng('rng_stream1')}")
-    # same value as second code snippet above
-    print(f"rng_stream2: {self.make_rng('rng_stream2')}")
-    # same value as first code snippet above
-    print(f"rng_stream1: {self.make_rng('rng_stream1')}")
-    # same value as second code snippet above
-    print(f"rng_stream2: {self.make_rng('rng_stream2')}")
-    # same value as first code snippet above
-    print(f"rng_stream1: {self.make_rng('rng_stream1')}")
-    # same value as second code snippet above
-    print(f"rng_stream2: {self.make_rng('rng_stream2')}")
-
-rng_module_two_streams = RNGModuleTwoStreams()
-variables = rng_module_two_streams.init(
-  {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(1)}
-)
-```
-
-Providing the same seed PRNG key will result in the same values being generated (provided that the same operations are used for those keys).
-
-```{code-cell}
-:outputId: 12e8c3c5-a32c-46ab-ba71-98cc90aaed70
-
-variables = rng_module_two_streams.init(
-  {'rng_stream1': jax.random.key(0), 'rng_stream2': jax.random.key(0)}
-)
-```
-
-### How `self.make_rng` works under the hood
-
-+++
-
-This is what happens when `self.make_rng` (`flax.linen.Module.make_rng`) is called:
-* The following data is collected:
-  * The path of the `Module` as provided by `self.scope.path` (the top-level root module has an empty path `()`).
-  * The `self.make_rng` call count. That is, the number of times `self.make_rng` has been called for this specific stream (including this call).
-    * **Note:** Each sub-`Module` will have its own individual call count that's separate from other `Module`s. For example, a `Module` that has called `self.make_rng('params')` twice and contains a sub-`Module` that has called `self.make_rng('params')` once, will have a call count of 2 and 1 for each of the RNG stream `'params'`, respectively.
-* The data is bundled into a tuple and fed into a hash function and produces an integer.
-* The generated integer is folded into the RNG stream's starting seed PRNG key to generate a new, unique PRNG key.
-
-+++
-
-Below is a slightly simplified version of the hash function that Flax uses for `self.make_rng`:
-
-```{code-cell}
-def produce_hash(data):
-  m = hashlib.sha1()
-  for x in data:
-    if isinstance(x, str):
-      m.update(x.encode('utf-8'))
-    elif isinstance(x, int):
-      m.update(x.to_bytes((x.bit_length() + 7) // 8, byteorder='big'))
-    else:
-      raise ValueError(f'Expected int or string, got: {x}')
-  d = m.digest()
-  hash_int = int.from_bytes(d[:4], byteorder='big')
-  return hash_int
-```
-
-And now you can manually reproduce the PRNG keys generated from `self.make_rng`:
-
-```{code-cell}
-:outputId: 6bf352c1-086a-4c7a-8ab9-92abde614270
-
-stream_seed = jax.random.key(0)
-for call_count in range(1, 4):
-  hash_int = produce_hash(data=(call_count,))
-  print(jax.random.fold_in(stream_seed, jnp.uint32(hash_int)))
-```
-
-```{code-cell}
-:outputId: 3809fddc-7319-4d55-ccc5-7b7c78160ea9
-
-variables = rng_module.init({'rng_stream': jax.random.key(0)})
-```
-
-### Sub-`Module`s and `self.make_rng`
-
-+++
-
-This section explores how `self.make_rng` (`flax.linen.Module.make_rng`) behaves with sub-`Module`s.
-
-Consider the following example:
-
-```{code-cell}
-:outputId: cc88c770-80ff-4093-e37f-6dc10c24006f
-
-class RNGSubSubModule(nn.Module):
-  def __call__(self):
-    print(f"{self.name}, count 1: {self.make_rng('rng_stream')}")
-    print(f"{self.name}, count 2: {self.make_rng('rng_stream')}")
-
-class RNGSubModule(nn.Module):
-  @nn.compact
-  def __call__(self):
-    print(f"{self.name}, count 1: {self.make_rng('rng_stream')}")
-    print(f"{self.name}, count 2: {self.make_rng('rng_stream')}")
-    RNGSubSubModule()()
-
-class RNGModule(nn.Module):
-  @nn.compact
-  def __call__(self):
-    print(f"RNGModule, count 1: {self.make_rng('rng_stream')}")
-    print(f"RNGModule, count 2: {self.make_rng('rng_stream')}")
-    RNGSubModule()()
-
-rng_module = RNGModule()
-variables = rng_module.init({'rng_stream': jax.random.key(0)})
-```
-
-As previously discussed, the data that is fed into the Flax hash function consists of:
-
-  * The path of the `Module`, provided by `self.scope.path` (the top-level root module has an empty path `()`); and
-  * The call count for the specific RNG stream.
-
-In addition, note that each Flax `Module` and sub-`Module` have their own individual call counts, even for the same RNG stream. The convention for sub-`Module` names is: `f'{module_name}_{module_number}'`. For example, the first `Dense` sub-`Module` will be called `Dense_0`, the second one will be called `Dense_1`, and so on.
-
-Therefore, the following data will be fed into the hash function:
-
-  * For `RNGModule`: The data is just the call count, such as `(1,)` and `(2,)`, since the root `Module` has an empty path.
-  * For `RNGSubModule`: The data is `('RNGSubModule_0', 1)` and `('RNGSubModule_0', 2)`.
-  * For `RNGSubSubModule`: The data is `('RNGSubModule_0', 'RNGSubSubModule_0', 1)` and `('RNGSubModule_0', 'RNGSubSubModule_0', 2)`.
-
-With this data, you can manually reproduce the PRNG keys generated from the `Module` and sub-`Module`s using `self.make_rng`.
-
-For example:
-
-```{code-cell}
-:outputId: d38896e6-8061-4f54-fa30-680b3f524071
-
-stream_seed = jax.random.key(0)
-for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_0', 'RNGSubSubModule_0')):
-  if initial_data:
-    module_name = initial_data[-1]
-  else:
-    module_name = 'RNGModule'
-  for call_count in (1, 2):
-    hash_int = produce_hash(data=initial_data+(call_count,))
-    rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))
-    print(f"{module_name}, count {call_count}: {rng_key}")
-```
-
-If the same sub-`Module` class is used multiple times, you can increment the suffix of the sub-`Module` name accordingly. For example: `RNGSubModule_0`, `RNGSubModule_1`, and so on.
-
-```{code-cell}
-:outputId: 556a79c6-0b38-4736-f5ad-79fd1976b191
-
-class RNGSubModule(nn.Module):
-  @nn.compact
-  def __call__(self):
-    print(f"{self.name}, count 1: {self.make_rng('rng_stream')}")
-    print(f"{self.name}, count 2: {self.make_rng('rng_stream')}")
-
-class RNGModule(nn.Module):
-  @nn.compact
-  def __call__(self):
-    print(f"RNGModule, count 1: {self.make_rng('rng_stream')}")
-    print(f"RNGModule, count 2: {self.make_rng('rng_stream')}")
-    RNGSubModule()()
-    RNGSubModule()()
-
-rng_module = RNGModule()
-variables = rng_module.init({'rng_stream': jax.random.key(0)})
-```
-
-```{code-cell}
-:outputId: 433f5e48-d379-490f-e499-ef4c24032776
-
-stream_seed = jax.random.key(0)
-for initial_data in ((), ('RNGSubModule_0',), ('RNGSubModule_1',)):
-  if initial_data:
-    module_name = initial_data[-1]
-  else:
-    module_name = 'RNGModule'
-  for call_count in (1, 2):
-    hash_int = produce_hash(data=initial_data+(call_count,))
-    rng_key = jax.random.fold_in(stream_seed, jnp.uint32(hash_int))
-    print(f"{module_name}, count {call_count}: {rng_key}")
-```
-
-### Using `self.param` and `self.variable`
-
-+++
-
-Flax users have the option of creating additional parameters and variables in their modules by using the [`self.param`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.param) and [`self.variable`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/module.html#flax.linen.Module.variable) `Module` methods. An `init_fn` argument must be passed to these methods so that it can generate the initial value of the parameter/variable. `self.make_rng` is commonly used implicitly or explicitly in this `init_fn`, since many initializer functions are stochastic in nature and require a PRNG key. See the full list of Flax initializers [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/initializers.html).
-
-There are a couple of differences between the two methods that the user should take note of:
-* `self.param` always creates a parameter in the `'params'` [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Params-parameters), whereas `self.variable` creates a variable in any [collection](https://flax.readthedocs.io/en/latest/glossary.html#term-Variable-collections) the user specifies
-* `self.param` will automatically call `self.make_rng('params')` and pass in the generated PRNG key implicitly to the `init_fn` of the parameter you instantiated (it will be passed in as the first argument), whereas users will have to manually specify what RNG stream to call `self.make_rng` on in the `init_fn` of `self.variable` (it could be `'params'` or something different).
-
-Below is an example using both `self.param` and `self.variable`:
-
-```{code-cell}
-:outputId: fc7df4d6-5d6e-4c3c-98f2-d5b0e9a8a45e
-
-class Model(nn.Module):
-  @nn.compact
-  def __call__(self, x):
-    # kernel will use 'params' seed, initial data will include 'Dense_0', call count 1
-    x = nn.Dense(2, kernel_init=jax.random.normal, use_bias=False)(x)
-    # model_param will use 'params' seed, call count 1
-    model_param = self.param('model_param', jax.random.normal, x.shape)
-    # model_variable1 will use 'params' seed, call count 2
-    model_variable1 = self.variable(
-      'other_collection',
-      'model_variable1',
-      lambda: jax.random.normal(self.make_rng('params'), x.shape),
-    )
-    # model_variable2 will use 'other' seed, call count 1
-    model_variable2 = self.variable(
-      'other_collection',
-      'model_variable2',
-      lambda: jax.random.normal(self.make_rng('other'), x.shape),
-    )
-    # kernel will use 'params' seed, initial data will include 'Dense_1', call count 1
-    # bias will use 'params' seed, initial data will include 'Dense_1', call count 2
-    x = nn.Dense(2, kernel_init=jax.random.normal, bias_init=jax.random.normal)(
-      x
-    )
-    return x
-
-model = Model()
-variables = model.init(
-  {'params': jax.random.key(0), 'other': jax.random.key(1)}, jnp.ones((2, 2))
-)
-print(variables['params']['Dense_0']['kernel'])
-print(variables['params']['model_param'])
-print(variables['other_collection']['model_variable1'])
-print(variables['other_collection']['model_variable2'])
-print(variables['params']['Dense_1']['kernel'])
-print(variables['params']['Dense_1']['bias'])
-```
-
-Remember:
-* there is a separate count for each RNG stream; this is why the count for `self.make_rng('other')` starts at 1 even though there were earlier calls of `self.make_rng('params')`
-* each submodule has their own separate count for each rng stream; this is why each `Dense` layer has their own separate count for `self.make_rng('params')` and why `model_param` and `model_variable1` share the same count (since they are defined within the same top-level parent module)
-
-```{code-cell}
-:outputId: 40745795-8bbc-4c31-811a-f3658e5459d7
-
-params_seed = jax.random.key(0)
-other_seed = jax.random.key(1)
-for initial_data, count, seed, shape in (
-  (('Dense_0',), 1, params_seed, (2, 2)),
-  ((), 1, params_seed, (2, 2)),
-  ((), 2, params_seed, (2, 2)),
-  ((), 1, other_seed, (2, 2)),
-  (('Dense_1',), 1, params_seed, (2, 2)),
-  (('Dense_1',), 2, params_seed, (1, 2)),
-):
-  hash_int = produce_hash(data=(*initial_data, count))
-  rng_key = jax.random.fold_in(seed, jnp.uint32(hash_int))
-  print(jax.random.normal(rng_key, shape))
-```
-
-### Managing RNG streams inside a training loop
-
-+++
-
-Below is an example of managing RNG streams from `self.make_rng`, `self.param`, `self.variable` and `nn.Dropout` in a training loop (note: `nn.Dropout` requires a seed PRNG key to be passed in the `'dropout'` RNG stream, since it implicitly calls `self.make_rng('dropout')`):
-
-```{code-cell}
-class SubModule(nn.Module):
-  @nn.compact
-  def __call__(self, x, train):
-    # variables created using `self.param` will use `self.make_rng('params')`
-    kernel = self.param('submodule_kernel', jax.random.normal, x.shape)
-    x = x + kernel
-    # `nn.Dropout` will use self.make_rng('dropout')
-    x = nn.Dropout(0.2)(x, deterministic=not train)
-    # `nn.Dense` will use self.make_rng('params')
-    x = nn.Dense(3)(x)
-    return x
-
-class Model(nn.Module):
-  @nn.compact
-  def __call__(self, x, train):
-    # make kernel use `self.make_rng('other')`
-    kernel = self.variable(
-      'other_collection',
-      'module_kernel',
-      lambda: jax.random.normal(self.make_rng('other'), x.shape),
-    )
-    x = (
-      x + kernel.value
-    )  # `.value` will extract the underlying value of the variable
-    x = SubModule()(x, train)
-    # `nn.Dropout` will use self.make_rng('dropout')
-    x = nn.Dropout(0.2)(x, deterministic=not train)
-    # `nn.Dense` will use self.make_rng('params')
-    x = nn.Dense(2)(x)
-    return x
-
-params_rng, other_rng, train_rng = jax.random.split(jax.random.key(0), 3)
-init_rngs = {'params': params_rng, 'other': other_rng}
-
-x = jnp.ones((1, 3))
-y = jnp.ones((1, 2))
-
-module = Model()
-variables = module.init(init_rngs, x, train=False)
-```
-
-```{code-cell}
-:outputId: 9e689bfd-ba73-4b78-c3c0-89693caada2d
-
-def update(variables, rng):
-  # we don't need to provide a 'params' or 'other' rng, as only 'dropout' rng will be used during training
-  # split the rng to get a dropout_rng to be used for this training iteration,
-  # and to get another rng key to be used for the next training iteration
-  dropout_rng, next_rng = jax.random.split(rng)
-  def loss(params):
-    out = module.apply(
-      {'params': params, 'other_collection': variables['other_collection']},
-      x,
-      train=True,
-      rngs={'dropout': dropout_rng},
-    )
-    return jnp.mean((y - out) ** 2)
-  grads = jax.grad(loss)(variables['params'])
-  params = jax.tree_map(lambda p, g: p - 1e-3 * g, variables['params'], grads)
-  return {
-    'params': params,
-    'other_collection': variables['other_collection'],
-  }, next_rng
-
-for _ in range(10):
-  variables, train_rng = update(variables, train_rng)
-  out = module.apply(variables, x, train=False)
-  print(jnp.mean((y - out)**2))
-```
-
-### 🔪 Sharp edge 🔪 - unintentionally generating the same values
-
-There is an edge case where the same value can be unintentionally generated.
-See the [Flax issue](https://github.com/google/flax/issues/2157) for more details.
-
-```{code-cell}
-:outputId: 80e45c76-bb4d-48ab-8fac-8b14126428d1
-
-class Leaf(nn.Module):
-  def __call__(self, x):
-    return x + jax.random.randint(self.make_rng("rng"), (), 0, 100)
-
-class Node(nn.Module):
-  leaf_name: str
-  @nn.compact
-  def __call__(self, x):
-    return Leaf(name=self.leaf_name)(x)
-
-class Model(nn.Module):
-  @nn.compact
-  def __call__(self, x):
-    return (Node(name="ab", leaf_name="cdef")(x),
-            Node(name="abc", leaf_name="def")(x),
-    )
-
-out1, out2 = Model().apply({}, 0, rngs={"rng": jax.random.key(33)})
-out1 == out2 # same output, despite having different submodule names
-```
-
-This occurs because the hash function [concatenates strings together](https://docs.python.org/3/library/hashlib.html#hashlib.hash.update), so the data `('AB', 'C')` is equivalent to data `('A', 'BC')` when fed into the hash function, therefore producing the same hash int.
-
-```{code-cell}
-:outputId: 07fc8349-41a3-4c74-d25c-0ffc4aa1be0a
-
-print(produce_hash(data=('A', 'B', 'C', 1)))
-print(produce_hash(data=('AB', 'C', 1)))
-print(produce_hash(data=('A', 'BC', 1)))
-print(produce_hash(data=('ABC', 1)))
-```
-
-To avoid this edge case, users can flip the `flax_fix_rng_separator` [configuration flag](https://flax.readthedocs.io/en/latest/api_reference/flax.config.html#flax.configurations.Config.flax_fix_rng_separator) to `True`.
-
-```{code-cell}
-:outputId: 4ba0c3cb-e903-40af-ab3f-687d83c257c9
-
-flax.config.update('flax_fix_rng_separator', True)
-out1, out2 = Model().apply({}, 0, rngs={"rng": jax.random.key(33)})
-out1 == out2 # different output
-```
-
-## Managing RNG's on multiple devices
-
-+++
-
-This section will show examples on how to use `jit` and `shard_map` to use RNG's in multi-device settings.
-
-+++
-
-### Using `jax.jit`
-
-+++
-
-When using [`jax.jit`](https://jax.readthedocs.io/en/latest/_autosummary/jax.jit.html), we can use RNG's as we did before, but we now include `in_shardings` and `out_shardings` arguments to shard specify how to shard input and output data.
-
-For more details on training on multiple devices in Flax using `jax.jit`, see our [Scale up Flax Modules on multiple devices guide](https://flax.readthedocs.io/en/latest/guides/parallel_training/flax_on_pjit.html#) and [lm1b example](https://github.com/google/flax/tree/main/examples/lm1b).
-
-```{code-cell}
-:outputId: 9deab9d8-3e15-4be0-8d91-279f6984ee99
-
-# Create a mesh and annotate the axis with a name.
-device_mesh = mesh_utils.create_device_mesh((8,))
-print(device_mesh)
-
-mesh = Mesh(devices=device_mesh, axis_names=('data',))
-print(mesh)
-
-data_sharding = NamedSharding(mesh, PartitionSpec('data',))
-print(data_sharding)
-```
-
-```{code-cell}
-:outputId: 9b47eef2-612e-4d9f-ffb2-1e868cb52d86
-
-class Model(nn.Module):
-  @nn.compact
-  def __call__(self, x, add_noise):
-    x = nn.Dense(1)(x)
-    # use jnp.where for control flow; for more details see: https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError
-    return jnp.where(
-      add_noise, x + jax.random.normal(self.make_rng('params'), x.shape), x
-    )
-
-module = Model()
-init_rng, apply_rng = jax.random.split(jax.random.key(0))
-x = jnp.ones((8, 1))
-variables = module.init(init_rng, x, False)
-
-# create custom forward function, since jit does not support kwargs when in_shardings is specified
-def forward(variables, x, add_noise, rng):
-  return module.apply(variables, x, add_noise, rngs={'params': rng})
-
-# shard the inputs x across devices
-# replicate the variables, add_noise boolean and rng key across devices
-# shard the output across devices
-jit_forward = jax.jit(
-  forward,
-  in_shardings=(None, data_sharding, None, None),
-  out_shardings=data_sharding,
-)
-out = jit_forward(variables, x, True, apply_rng)
-out
-```
-
-The output is different given the same input, meaning the RNG key was used to add noise to the output.
-
-We can also confirm that the output is sharded across devices:
-
-```{code-cell}
-:outputId: f6596bdc-89b4-46bf-ba99-84b1009d8156
-
-out.addressable_shards
-```
-
-Another way to visualize the output sharding:
-
-```{code-cell}
-:outputId: 7df87ccd-0979-4b85-fbac-1d31aac53276
-
-jax.debug.visualize_array_sharding(out)
-```
-
-If we choose not to add noise, then the output is the same across all batches (as expected, since the input is the same for all batches):
-
-```{code-cell}
-:outputId: cb7eb7d0-3fa7-4f72-a129-18901dc61bb1
-
-out = jit_forward(variables, x, False, apply_rng)
-print(out)
-```
-
-We can confirm the un-jitted function produces the same values, albeit unsharded (note there may be small numerical differences due to compiler optimizations from jitting):
-
-```{code-cell}
-:outputId: b88834c1-ebb2-422a-a665-1732015a3974
-
-out = forward(variables, x, True, apply_rng)
-out
-```
-
-```{code-cell}
-:outputId: 5b06f58b-7f1f-4cd6-c29f-3da5c0ea189f
-
-out = forward(variables, x, False, apply_rng)
-out
-```
-
-### Using `shard_map`
-
-+++
-
-When using [`jax.experimental.shard_map.shard_map`](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html), the important parts to remember are to:
-* split your PRNG key to produce a different key for each device
-* the PRNG keys will be sharded automatically to each device (provided you use the correct partition specification), but the [**rank of the original batched PRNG key array will not be reduced**](https://jax.readthedocs.io/en/latest/jep/14273-shard-map.html#rank-reducing-vs-rank-preserving-maps-over-array-axes); e.g.
-with a batch of 8 PRNG keys and 8 devices, each device will see a PRNG key batch of size 1 within the `shard_map`-ed function
-  * therefore to access the PRNG key itself, we need to index slice into it (see the example below)
-
-```{code-cell}
-:outputId: 29b2911f-28a7-4a25-fcf7-4c3544dafada
-
-def forward(variables, x, add_noise, rng_key_batch):
-  # rng_key_batch is a batch of size 1 containing 1 PRNG key
-  # index slice into the rng_key_batch to access the PRNG key
-  return module.apply(
-    variables, x, add_noise, rngs={'params': rng_key_batch[0]}
-  )
-
-# define partition specifications
-data_pspec = PartitionSpec('data')
-no_pspec = PartitionSpec()
-
-# shard the inputs x and rng keys across devices
-# replicate the variables and add_noise boolean across devices
-# shard the output across devices
-shmap_forward = shard_map(
-  forward,
-  mesh=mesh,
-  in_specs=(no_pspec, data_pspec, no_pspec, data_pspec),
-  out_specs=data_pspec,
-)
-# get 8 different rng's that will be used by the 8 devices when doing forward inference
-apply_rngs = jax.random.split(apply_rng, 8)
-out = shmap_forward(variables, x, True, apply_rngs)
-out
-```
-
-Confirm that the output is sharded across devices:
-
-```{code-cell}
-:outputId: 2d74dff9-3a7d-47df-afbc-3bda4b8edf51
-
-out.addressable_shards
-```
-
-```{code-cell}
-:outputId: 28bfe59b-ded1-455a-c894-f1c457ec28bf
-
-jax.debug.visualize_array_sharding(out)
-```
-
-## Lifted transforms
-
-+++
-
-[Flax lifted transforms](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/transformations.html) allow you to use [JAX transforms](https://github.com/google/jax#transformations) with `Module` arguments. This section will show you how to control how PRNG keys are split in Flax lifted transforms.
-
-Refer to [Lifted transformations](https://flax.readthedocs.io/en/latest/developer_notes/lift.html) for more detail.
-
-+++
-
-### `nn.vmap`
-
-+++
-
-We can use [`nn.vmap`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.vmap.html) to create a batched `Dense` layer:
-
-```{code-cell}
-:outputId: 7d7c4c3c-2ed2-40dd-f8d3-bb1ef91d93eb
-
-x = jnp.ones((3, 2))
-
-BatchDense = nn.vmap(
-    nn.Dense,
-    in_axes=0, out_axes=0,
-    variable_axes={'params': None},
-    split_rngs={'params': False})
-
-BatchDense(2).init(jax.random.key(0), x)
-```
-
-By denoting `variable_axes={'params': 0}'`, we vectorize the `params` Arrays on the first axis. However the parameter values generated are all identical to each other:
-
-```{code-cell}
-:outputId: 93152c8f-ada9-4cf5-c0e2-560a284b3981
-
-BatchDense = nn.vmap(
-    nn.Dense,
-    in_axes=0, out_axes=0,
-    variable_axes={'params': 0},
-    split_rngs={'params': False})
-
-BatchDense(2).init(jax.random.key(0), x)
-```
-
-If we also make `split_rngs={'params': True}`, then the PRNG key we provide is split across the variable axis (in this case, the batch axis 0), and we can generate different parameters for each batch input:
-
-```{code-cell}
-:outputId: dfbf4e49-55dc-4a6f-bece-396cb48d875e
-
-BatchDense = nn.vmap(
-    nn.Dense,
-    in_axes=0, out_axes=0,
-    variable_axes={'params': 0},
-    split_rngs={'params': True})
-
-BatchDense(2).init(jax.random.key(0), x)
-```
-
-Adding a variable via `self.variable` is straightforward:
-
-```{code-cell}
-:outputId: 20fa39a3-9e7a-4a58-8ad7-48a0119ec466
-
-class Model(nn.Module):
-  @nn.compact
-  def __call__(self, x):
-    x = nn.Dense(2)(x)
-    kernel = self.variable(
-      'other_collection',
-      'kernel',
-      lambda: jax.random.normal(self.make_rng('other'), x.shape),
-    )
-    return x + kernel.value
-
-BatchModel = nn.vmap(
-  Model,
-  in_axes=0,
-  out_axes=0,
-  variable_axes={'params': 0, 'other_collection': 0},
-  split_rngs={'params': True, 'other': True},
-)
-
-BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)
-```
-
-We can control which RNG stream to split, for example, if we only wanted to split the `'params'` RNG stream, then the variables generated from `self.variable` will be the same for each batch input:
-
-```{code-cell}
-:outputId: 69d52ec5-9018-43c8-9d77-65e9bf00b79f
-
-BatchModel = nn.vmap(
-    Model,
-    in_axes=0, out_axes=0,
-    variable_axes={'params': 0, 'other_collection': 0},
-    split_rngs={'params': True, 'other': False})
-
-BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)
-```
-
-We can also control which parameters / variables should be generated for each batch input, for example, if we only wanted `'params'` to generate separate parameters for each batch input:
-
-```{code-cell}
-:outputId: b01187f9-aded-41bc-d17f-e2965355675a
-
-BatchModel = nn.vmap(
-    Model,
-    in_axes=0, out_axes=0,
-    variable_axes={'params': 0, 'other_collection': None},
-    split_rngs={'params': True, 'other': False})
-
-BatchModel().init({'params': jax.random.key(0), 'other': jax.random.key(1)}, x)
-```
-
-### `nn.scan`
-
-+++
-
-We can use [`nn.scan`](https://flax.readthedocs.io/en/latest/api_reference/flax.linen/_autosummary/flax.linen.scan.html) to create a scanned `Module` layer (this is useful for simplifying repetitively stacked submodules):
-
-```{code-cell}
-:outputId: fda568b4-3c1d-4a5e-96f0-058c6fc5b49a
-
-x = jnp.ones((3, 2))
-
-class ResidualMLPBlock(nn.Module):
-  @nn.compact
-  def __call__(self, x, _):
-    h = nn.Dense(features=2)(x)
-    h = nn.relu(h)
-    return x + h, None # return an empty carry
-
-ScanMLP = nn.scan(
-      ResidualMLPBlock, variable_axes={'params': 0},
-      variable_broadcast=False, split_rngs={'params': True},
-      length=3)
-
-ScanMLP().init(jax.random.key(0), x, None) # pass in an empty carry
-```
-
-Similar to before, we can control whether to split the RNG stream or not, for example, if we wanted all the stacked modules to be initialized to the same parameter values, we can pass in `split_rngs={'params': False}`:
-
-```{code-cell}
-:outputId: 620aff62-e20d-4794-b3af-5a5058c2471d
-
-ScanMLP = nn.scan(
-      ResidualMLPBlock, variable_axes={'params': 0},
-      variable_broadcast=False, split_rngs={'params': False},
-      length=3)
-
-ScanMLP().init(jax.random.key(0), x, None)
-```
diff --git a/flax/core/meta.py b/flax/core/meta.py
index 007825e1..92c4312f 100644
--- a/flax/core/meta.py
+++ b/flax/core/meta.py
@@ -293,10 +293,9 @@ def with_partitioning(
 
   Example::
 
-    >>> import flax.linen as nn
-    >>> kernel_init = nn.with_partitioning(
-    ...     nn.initializers.lecun_normal(), (None, "data"))
-    >>> partitioned_dense = nn.Dense(features=3, kernel_init=kernel_init)
+    kernel_init = with_partitioning(
+        nn.initializers.lecun_normal, (None, "data"))
+    partitioned_dense = nn.Dense(features, kernel_init=kernel_init)
 
   Args:
     fn: The function to be wrapped. Typically this is an initializer.
diff --git a/flax/linen/__init__.py b/flax/linen/__init__.py
index 3106546c..d78f3881 100644
--- a/flax/linen/__init__.py
+++ b/flax/linen/__init__.py
@@ -118,6 +118,7 @@ from .module import (
 from .normalization import (
     BatchNorm as BatchNorm,
     GroupNorm as GroupNorm,
+    InstanceNorm as InstanceNorm,
     LayerNorm as LayerNorm,
     RMSNorm as RMSNorm,
     SpectralNorm as SpectralNorm,
diff --git a/flax/linen/combinators.py b/flax/linen/combinators.py
index cf566dba..b20195ad 100644
--- a/flax/linen/combinators.py
+++ b/flax/linen/combinators.py
@@ -16,7 +16,7 @@
 
 from typing import Any, Callable, Dict, Sequence
 
-from flax.linen.module import Module, compact
+from flax.linen.module import Module
 
 
 class Sequential(Module):
@@ -43,15 +43,6 @@ class Sequential(Module):
     ...                           nn.Dense(2),
     ...                           nn.log_softmax])(x)
 
-  Since `Sequential.__call__` is a `compact` method, you can also pass functions
-  that construct Modules inline if you need shape inference::
-
-    module = nn.Sequential([
-        # << more layers
-        lambda x: SomeModule(x.shape[-1])(x), # shape inference
-        # << more layers
-    ])
-
   This combinator supports also layers that return multiple outputs if returned
   as a tuple or a dictionary. If the output of a layer is a ``tuple`` it will be
   expanded as ``*args`` in the next layer, if its a ``dict`` it
@@ -79,13 +70,6 @@ class Sequential(Module):
     ...   def __call__(self, x):
     ...     return nn.Sequential([CrossAttentionBlock() for _ in
     ...                           range(self.num_layers)])(query, key_value)
-
-
-  Attributes:
-    layers: A sequence of callables to be applied in order.
-
-  Raises:
-    ValueError: If layers is not a sequence.
   """
 
   layers: Sequence[Callable[..., Any]]
@@ -97,7 +81,6 @@ class Sequential(Module):
       )
     super().__post_init__()
 
-  @compact
   def __call__(self, *args, **kwargs):
     if not self.layers:
       raise ValueError(f'Empty Sequential module {self.name}.')
diff --git a/flax/linen/normalization.py b/flax/linen/normalization.py
index cdf96ae6..1535a7e8 100644
--- a/flax/linen/normalization.py
+++ b/flax/linen/normalization.py
@@ -730,6 +730,127 @@ class GroupNorm(Module):
     )
 
 
+class InstanceNorm(Module):
+  """Instance normalization (https://arxiv.org/abs/1607.08022v3).
+
+  InstanceNorm normalizes the activations of the layer for each channel (rather
+  than across all channels like Layer Normalization), and for each given example
+  in a batch independently (rather than across an entire batch like Batch
+  Normalization). i.e. applies a transformation that maintains the mean activation
+  within each channel within each example close to 0 and the activation standard
+  deviation close to 1.
+
+  NOTE: This normalization operation is identical to LayerNorm and GroupNorm; the
+  difference is simply which axes are reduced and the shape of the feature axes
+  (i.e. the shape of the learnable scale and bias parameters).
+
+  Example usage::
+
+    >>> import flax.linen as nn
+    >>> import jax
+    >>> import numpy as np
+
+    >>> # dimensions: (batch, height, width, channel)
+    >>> x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))
+    >>> layer = nn.InstanceNorm()
+    >>> variables = layer.init(jax.random.key(1), x)
+    >>> variables
+    {'params': {'scale': Array([1., 1., 1., 1., 1.], dtype=float32), 'bias': Array([0., 0., 0., 0., 0.], dtype=float32)}}
+    >>> y = layer.apply(variables, x)
+
+    >>> # having a channel_axis of -1 in InstanceNorm is identical to reducing all non-batch,
+    >>> # non-channel axes and using the channel_axes as the feature_axes in LayerNorm
+    >>> y2 = nn.LayerNorm(reduction_axes=[1, 2], feature_axes=-1).apply(variables, x)
+    >>> np.testing.assert_allclose(y, y2, atol=1e-7)
+    >>> y3 = nn.GroupNorm(num_groups=x.shape[-1]).apply(variables, x)
+    >>> np.testing.assert_allclose(y, y3, atol=1e-7)
+
+  Attributes:
+    epsilon: A small float added to variance to avoid dividing by zero.
+    dtype: the dtype of the result (default: infer from input and params).
+    param_dtype: the dtype passed to parameter initializers (default: float32).
+    use_bias:  If True, bias (beta) is added.
+    use_scale: If True, multiply by scale (gamma). When the next layer is linear
+      (also e.g. nn.relu), this can be disabled since the scaling will be done
+      by the next layer.
+    bias_init: Initializer for bias, by default, zero.
+    scale_init: Initializer for scale, by default, one.
+    channel_axes: Axes for channel. This is considered the feature axes for the
+      learned bias and scaling parameter. All other axes except the batch axes
+      (which is assumed to be the leading axis) will be reduced.
+    axis_name: the axis name used to combine batch statistics from multiple
+      devices. See ``jax.pmap`` for a description of axis names (default: None).
+      This is only needed if the model is subdivided across devices, i.e. the
+      array being normalized is sharded across devices within a pmap or shard
+      map. For SPMD jit, you do not need to manually synchronize. Just make sure
+      that the axes are correctly annotated and XLA:SPMD will insert the
+      necessary collectives.
+    axis_index_groups: groups of axis indices within that named axis
+      representing subsets of devices to reduce over (default: None). For
+      example, ``[[0, 1], [2, 3]]`` would independently batch-normalize over the
+      examples on the first two and last two devices. See ``jax.lax.psum`` for
+      more details.
+    use_fast_variance: If true, use a faster, but less numerically stable,
+      calculation for the variance.
+  """
+
+  epsilon: float = 1e-6
+  dtype: Optional[Dtype] = None
+  param_dtype: Dtype = jnp.float32
+  use_bias: bool = True
+  use_scale: bool = True
+  bias_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.zeros
+  scale_init: Callable[[PRNGKey, Shape, Dtype], Array] = initializers.ones
+  channel_axes: Axes = -1
+  axis_name: Optional[str] = None
+  axis_index_groups: Any = None
+  use_fast_variance: bool = True
+
+  @compact
+  def __call__(self, x, mask=None):
+    """Applies instance normalization on the input.
+
+    Args:
+      x: the inputs
+      mask: Binary array of shape broadcastable to ``inputs`` tensor, indicating
+        the positions for which the mean and variance should be computed.
+
+    Returns:
+      Normalized inputs (the same shape as inputs).
+    """
+    channel_axes = _canonicalize_axes(x.ndim, self.channel_axes)
+    if 0 in channel_axes:
+      raise ValueError('The channel axes cannot include the leading dimension '
+                       'as this is assumed to be the batch axis.')
+    reduction_axes = [i for i in range(1, x.ndim) if i not in channel_axes]
+
+    mean, var = _compute_stats(
+      x,
+      reduction_axes,
+      self.dtype,
+      self.axis_name,
+      self.axis_index_groups,
+      use_fast_variance=self.use_fast_variance,
+      mask=mask,
+    )
+
+    return _normalize(
+      self,
+      x,
+      mean,
+      var,
+      reduction_axes,
+      channel_axes,
+      self.dtype,
+      self.param_dtype,
+      self.epsilon,
+      self.use_bias,
+      self.use_scale,
+      self.bias_init,
+      self.scale_init,
+    )
+
+
 class SpectralNorm(Module):
   """Spectral normalization. See:
 
diff --git a/flax/linen/spmd.py b/flax/linen/spmd.py
index 702164a3..77414b72 100644
--- a/flax/linen/spmd.py
+++ b/flax/linen/spmd.py
@@ -329,7 +329,7 @@ def with_logical_partitioning(
 
     >>> import flax.linen as nn
     >>> kernel_init = nn.with_logical_partitioning(
-    ...     nn.initializers.lecun_normal(), (None, "data"))
+    ...     nn.initializers.lecun_normal, (None, "data"))
     >>> partitioned_dense = nn.Dense(features=3, kernel_init=kernel_init)
 
   Args:
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
index 4aadf57b..07bb1dcb 100644
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -191,8 +191,7 @@ def get_module_scopes(module, args=None, kwargs=None):
       for f in dataclasses.fields(module)
       if f.name != 'parent' and f.init
     }
-    for leaf in jax.tree_util.tree_leaves(attrs):
-      get_scopes_inner(leaf)
+    jax.tree_util.tree_map(get_scopes_inner, attrs)
     scopes.append(module.scope)
 
   get_scopes(module)
diff --git a/flax/struct.py b/flax/struct.py
index 313a6ee6..c538adb6 100644
--- a/flax/struct.py
+++ b/flax/struct.py
@@ -32,7 +32,7 @@ def field(pytree_node=True, **kwargs):
 
 
 @dataclass_transform(field_specifiers=(field,))  # type: ignore[literal-required]
-def dataclass(clz: _T, **kwargs) -> _T:
+def dataclass(clz: _T) -> _T:
   """Create a class which can be passed to functional transformations.
 
   NOTE: Inherit from ``PyTreeNode`` instead to avoid type checking issues when
@@ -104,9 +104,7 @@ def dataclass(clz: _T, **kwargs) -> _T:
   if '_flax_dataclass' in clz.__dict__:
     return clz
 
-  if 'frozen' not in kwargs.keys():
-    kwargs['frozen'] = True
-  data_clz = dataclasses.dataclass(**kwargs)(clz)  # type: ignore
+  data_clz = dataclasses.dataclass(frozen=True)(clz)  # type: ignore
   meta_fields = []
   data_fields = []
   for field_info in dataclasses.fields(data_clz):


Test patch:
diff --git a/tests/linen/linen_combinators_test.py b/tests/linen/linen_combinators_test.py
index 4b16a77a..c9316c4a 100644
--- a/tests/linen/linen_combinators_test.py
+++ b/tests/linen/linen_combinators_test.py
@@ -170,22 +170,6 @@ class SequentialTest(absltest.TestCase):
     np.testing.assert_equal(out_query.shape, (3, 5))
     np.testing.assert_equal(out_key_value.shape, (9, 5))
 
-  def test_sequential_compact(self):
-
-    mlp = nn.Sequential([
-        lambda x: nn.Dense(x.shape[-1])(x),
-        nn.relu,
-        lambda x: nn.Dense(x.shape[-1])(x),
-        nn.relu,
-        lambda x: nn.Dense(x.shape[-1])(x),
-    ])
-
-    params = mlp.init(random.key(0), jnp.ones((3, 5)))['params']
-
-    self.assertIn('Dense_0', params)
-    self.assertIn('Dense_1', params)
-    self.assertIn('Dense_2', params)
-
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/linen/linen_test.py b/tests/linen/linen_test.py
index 17b5f91c..88735070 100644
--- a/tests/linen/linen_test.py
+++ b/tests/linen/linen_test.py
@@ -430,6 +430,61 @@ class NormalizationTest(parameterized.TestCase):
     (y1, y2), _ = model.init_with_output(key, x)
     np.testing.assert_allclose(y1, y2, rtol=1e-4)
 
+  @parameterized.parameters(
+    {'channel_axes': -1},
+    {'channel_axes': (1, 2)},
+    {'channel_axes': (1, 2, 3)},
+    {'channel_axes': -1, 'use_fast_variance': False},
+  )
+  def test_instance_norm(self, channel_axes, use_fast_variance=True):
+    rng = random.key(0)
+    key1, key2 = random.split(rng)
+    e = 1e-5
+    x = random.normal(key1, (2, 3, 4, 5))
+    if not use_fast_variance:
+      x += 1e4  # This blows up fast variance, but should work otherwise.
+    model_cls = nn.InstanceNorm(
+      use_bias=False,
+      use_scale=False,
+      epsilon=e,
+      channel_axes=channel_axes,
+      use_fast_variance=use_fast_variance,
+    )
+    y, _ = model_cls.init_with_output(key2, x)
+    self.assertEqual(x.dtype, y.dtype)
+    self.assertEqual(x.shape, y.shape)
+
+    canonicalized_channel_axes = [
+      i if i >= 0 else (x.ndim + i)
+      for i in (
+        channel_axes if isinstance(channel_axes, tuple) else (channel_axes,)
+      )
+    ]
+    reduction_axes = [
+      i for i in range(1, x.ndim) if i not in canonicalized_channel_axes
+    ]
+    y_one_liner = (
+      x - x.mean(axis=reduction_axes, keepdims=True)
+    ) * jax.lax.rsqrt(x.var(axis=reduction_axes, keepdims=True) + e)
+
+    np.testing.assert_allclose(y_one_liner, y, atol=1e-6)
+
+  @parameterized.parameters(
+    {'channel_axes': 0},
+    {'channel_axes': -4},
+    {'channel_axes': (0, 3)},
+    {'channel_axes': (2, -4)},
+  )
+  def test_instance_norm_raise_error(self, channel_axes):
+    with self.assertRaisesRegex(
+      ValueError,
+      'The channel axes cannot include the leading dimension '
+      'as this is assumed to be the batch axis.',
+    ):
+      x = jax.random.normal(jax.random.key(0), (2, 3, 4, 5))
+      layer = nn.InstanceNorm(channel_axes=channel_axes)
+      _ = layer.init(jax.random.key(1), x)
+
   @parameterized.parameters(
     {
       'model_index': 0,
diff --git a/tests/struct_test.py b/tests/struct_test.py
index bfb3bd7e..efd9529a 100644
--- a/tests/struct_test.py
+++ b/tests/struct_test.py
@@ -47,21 +47,6 @@ class StructTest(absltest.TestCase):
     with self.assertRaises(dataclasses.FrozenInstanceError):
       p.y = 3
 
-  # Slots param is not supported in Python 3.9
-  # TODO(lew): Uncomment when Flax upgrades to Python 3.10.
-
-  # def test_slots(self):
-  #   slots_dataclass = partial(struct.dataclass, frozen=False, slots=True)
-  #   @slots_dataclass
-  #   class SlotsPoint:
-  #     x: float
-  #     y: float
-  #   p = SlotsPoint(x=1., y=2.)
-  #   p.x = 3.  # can assign to existing fields
-  #   self.assertEqual(p, SlotsPoint(x=3., y=2.))
-  #   with self.assertRaises(AttributeError):
-  #     p.z = 0.  # can't create new fields by accident.
-
   def test_pytree_nodes(self):
     p = Point(x=1, y=2, meta={'abc': True})
     leaves = jax.tree_util.tree_leaves(p)