diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
index 355334218..b645fea27 100644
--- a/.github/ISSUE_TEMPLATE/bug_report.md
+++ b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -10,11 +10,11 @@ assignees: ''
 Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.
 
 ### System information
-- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): 
-- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`: 
-- Python version: 
-- GPU/TPU model and memory: 
-- CUDA version (if applicable): 
+- OS Platform and Distribution (e.g., Linux Ubuntu 16.04):
+- Flax, jax, jaxlib versions (obtain with `pip show flax jax jaxlib`:
+- Python version:
+- GPU/TPU model and memory:
+- CUDA version (if applicable):
 
 
 ### Problem you have encountered:
diff --git a/.github/analytics/get_repo_metrics.py b/.github/analytics/get_repo_metrics.py
index 8ebce468a..319ad8b09 100644
--- a/.github/analytics/get_repo_metrics.py
+++ b/.github/analytics/get_repo_metrics.py
@@ -110,7 +110,7 @@ def parse_single_query(data, query_type):
   Parses the data returned by `send_query`
 
   .. warning::
-      
+
     Like `send_query`, the logic here depends on the specific structure
     of the query (e.g. it must be an issue or PR query, and must have a
     total count).
@@ -206,7 +206,7 @@ def _get_pr_features(prs):
     time_labeled_or_assigned = None
     time_merged_or_closed = None
     time_review = None
-    
+
     if pr["reviews"]["nodes"]:
       review = pr["reviews"]["nodes"][0]
       time_review = _to_datetime(review["createdAt"])
@@ -258,9 +258,9 @@ def _shift_n_months(date: datetime, n: int) -> datetime:
 
 
 def _rolling_window(
-  df: pd.DataFrame, 
-  f: Callable[[pd.DataFrame], pd.Series], 
-  window_size: int = 6, 
+  df: pd.DataFrame,
+  f: Callable[[pd.DataFrame], pd.Series],
+  window_size: int = 6,
   step: int = 1,
 ) -> pd.DataFrame:
   # start of month of the first issue
diff --git a/.github/analytics/issue_activity_since_date.gql b/.github/analytics/issue_activity_since_date.gql
index d23af57c9..5fa9a098f 100644
--- a/.github/analytics/issue_activity_since_date.gql
+++ b/.github/analytics/issue_activity_since_date.gql
@@ -1,7 +1,7 @@
 {
   # Queries all the issues in a repo. For each issue, we get some basic data such as
-  # the number, state, labels, and title. The most important part is the 'timelineItems' 
-  # which are the events that happened to the issue, we can use the information about 
+  # the number, state, labels, and title. The most important part is the 'timelineItems'
+  # which are the events that happened to the issue, we can use the information about
   # the datetime about certain key events to define some metrics. Note that we are
   # getting more information than is probably needed but its fine for now.
   repository(owner: "_REPO_OWNER_", name: "_REPO_NAME_") {
@@ -52,7 +52,7 @@
                 }
               }
             }
-          } 
+          }
         }
       }
     }
diff --git a/.github/analytics/pr_data_query.gql b/.github/analytics/pr_data_query.gql
index c8a6de18e..867e0934f 100644
--- a/.github/analytics/pr_data_query.gql
+++ b/.github/analytics/pr_data_query.gql
@@ -1,9 +1,9 @@
 query {
   # Queries all the Pull Requests in a repo. For each issue, we get some basic data such as
-  # the number, state, reviews, and title. The most important part is the 'timelineItems' 
-  # which are the events that happened to the issue, we can use the information about 
+  # the number, state, reviews, and title. The most important part is the 'timelineItems'
+  # which are the events that happened to the issue, we can use the information about
   # the datetime about certain key events to define some metrics. We also use the 'reviews'
-  # as indicators for certain metrics.  Note that we are getting more information than is 
+  # as indicators for certain metrics.  Note that we are getting more information than is
   # probably needed but its fine for now.
   repository(owner:"_REPO_OWNER_", name:"_REPO_NAME_") {
     pullRequests(first:100) {
diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index 0cba729fa..dbc4a5a08 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -2,7 +2,7 @@
 
 <!--
 
-Great, you are contributing to Flax! 
+Great, you are contributing to Flax!
 
 But... please read the following carefully so we can make sure your PR is merged
 easily.
diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 7e2c51d32..5013dec14 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -36,14 +36,14 @@ jobs:
     - name: Cancel previous
       uses: styfle/cancel-workflow-action@0.8.0
       with:
-        access_token: ${{ github.token }}          
+        access_token: ${{ github.token }}
     - uses: actions/checkout@v2
     # We allow at most 5 commits in a branch to ensure our CI doesn't break.
     - name: Check commit count in PR
       if: always()
       shell: bash
       run: |
-        set -x 
+        set -x
         # $GITHUB_REF is in format `refs/heads/<branch_name>`. We fetch it under
         # the name `commit-count` so we can refer to it below.
         # Do an unshallow fetch so we retrieve all commits (this is necessary
@@ -75,7 +75,7 @@ jobs:
     - name: Cancel previous
       uses: styfle/cancel-workflow-action@0.8.0
       with:
-        access_token: ${{ github.token }}          
+        access_token: ${{ github.token }}
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v1
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c5fc8f7cb..338cf38b1 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -53,7 +53,7 @@ New features:
 - Added async option to `save_checkpoints()` on single-process scenario.
 - Improved documentation pages.
 
-Bug fixes: 
+Bug fixes:
 - Fixed variable aliasing in put_variable
 - Fixed missing passthrough of nn.scan unroll arg
 - Fixed the MNIST example
diff --git a/README.md b/README.md
index c3f411ee5..f34698ed6 100644
--- a/README.md
+++ b/README.md
@@ -87,7 +87,7 @@ To upgrade to the latest version of Flax, you can use:
 
 ## What does Flax look like?
 
-We provide three examples using the Flax API: a simple multi-layer perceptron, a CNN and an auto-encoder. 
+We provide three examples using the Flax API: a simple multi-layer perceptron, a CNN and an auto-encoder.
 
 To learn more about the `Module` abstraction, see our [docs](https://flax.readthedocs.io/), our [broad intro to the Module abstraction](https://github.com/google/flax/blob/main/docs/notebooks/linen_intro.ipynb). For additional concrete demonstrations of best practices, see our
 [HOWTO guides](https://flax.readthedocs.io/en/latest/howtos.html).
@@ -174,11 +174,11 @@ decoded = model.apply(variables, encoded, method=model.decode)
 
 ## 🤗 Hugging Face
 
-In-detail examples to train and evaluate a variety of Flax models for 
-Natural Language Processing, Computer Vision, and Speech Recognition are 
+In-detail examples to train and evaluate a variety of Flax models for
+Natural Language Processing, Computer Vision, and Speech Recognition are
 actively maintained in the [🤗 Transformers repository](https://github.com/huggingface/transformers/tree/master/examples/flax).
 
-As of October 2021, the [19 most-used Transformer architectures](https://huggingface.co/transformers/#supported-frameworks) are supported in Flax 
+As of October 2021, the [19 most-used Transformer architectures](https://huggingface.co/transformers/#supported-frameworks) are supported in Flax
 and over 5000 pretrained checkpoints in Flax have been uploaded to the [🤗 Hub](https://huggingface.co/models?library=jax&sort=downloads).
 
 ## Citing Flax
diff --git a/dev/.devcontainer/devcontainer.json b/dev/.devcontainer/devcontainer.json
index c8f99a708..32b87efa7 100644
--- a/dev/.devcontainer/devcontainer.json
+++ b/dev/.devcontainer/devcontainer.json
@@ -5,9 +5,9 @@
 	"context": "..",
 	"dockerFile": "Dockerfile",
 
-	// Use 'settings' to set *default* container specific settings.json values on container create. 
+	// Use 'settings' to set *default* container specific settings.json values on container create.
 	// You can edit these settings after create using File > Preferences > Settings > Remote.
-	"settings": { 
+	"settings": {
 		"terminal.integrated.shell.linux": "/bin/bash",
 		"python.pythonPath": "/usr/local/bin/python",
 		"python.linting.enabled": true,
@@ -24,7 +24,7 @@
 	"runArgs": ["-v", "${env:HOME}${env:USERPROFILE}/.ssh:/root/.ssh-localhost:ro"],
 	"postCreateCommand": "sudo cp -r /root/.ssh-localhost ~/.ssh && sudo chown -R $(id -u):$(id -g) ~/.ssh && chmod 700 ~/.ssh && chmod 600 ~/.ssh/* && pip install -e ./flax",
 
-	// Uncomment the next line to have VS Code connect as an existing non-root user in the container. 
+	// Uncomment the next line to have VS Code connect as an existing non-root user in the container.
 	// On Linux, by default, the container user's UID/GID will be updated to match your local user. See
 	// https://aka.ms/vscode-remote/containers/non-root for details on adding a non-root user if none exist.
 	// "remoteUser": "vscode",
diff --git a/docs/README.md b/docs/README.md
index ebe546ed6..07e1dd7b1 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -45,7 +45,7 @@ use the following notation:
 #
 #   def sum(a, b):
 #     return a + b
-#   
+#
 #   sum(0, 1)
 ```
 
@@ -87,19 +87,19 @@ You can use the second method if the `path_to_obj` is very long. Some examples:
 # :func:`my_func`
 
 # Create a reference "Module.apply()" to method flax.linen.Module.apply.
-# :meth:`Module.apply() <flax.linen.Module.apply>`  # 
-``` 
+# :meth:`Module.apply() <flax.linen.Module.apply>`  #
+```
 
 To creata a hyperlink, use the following syntax:
 ```bash
-# Note the double underscore at the end: 
+# Note the double underscore at the end:
 # `Link to Google <http://www.google.com>`__
 ```
 
 ### How to specify arguments for classes and methods
 
 *  Class attributes should be specified using the `Attributes:` tag.
-*  Method argument should be specified using the `Args:` tags. 
+*  Method argument should be specified using the `Args:` tags.
 *  All attributes and arguments should have types.
 
 Here is an example from our library:
diff --git a/docs/_ext/codediff.py b/docs/_ext/codediff.py
index b6cc22975..9bccb3bf3 100644
--- a/docs/_ext/codediff.py
+++ b/docs/_ext/codediff.py
@@ -75,10 +75,10 @@ def _code_block(self, lines):
 
   def _tabs(self, *contents: Tuple[str, List[str]], sync):
     output = ['.. tab-set::'] + ['  ']
-    
+
     for title, content in contents:
       output += [f'  .. tab-item:: {title}']
-      
+
       if sync:
         key = title.strip()
         output += [f'    :sync: {key}']
diff --git a/docs/_ext/codediff_test.py b/docs/_ext/codediff_test.py
index 46d9cdf2a..9027d9491 100644
--- a/docs/_ext/codediff_test.py
+++ b/docs/_ext/codediff_test.py
@@ -24,7 +24,7 @@ class CodeDiffTest(absltest.TestCase):
   def test_parse(self):
 
     input_text = r'''@jax.jit #!
-def get_initial_params(key):   #!  
+def get_initial_params(key):   #!
   init_val = jnp.ones((1, 28, 28, 1), jnp.float32)
   initial_params = CNN().init(key, init_val)['params']
   extra_line
diff --git a/docs/_templates/autosummary/flax_module.rst b/docs/_templates/autosummary/flax_module.rst
index dd0031ac8..5f51933cf 100644
--- a/docs/_templates/autosummary/flax_module.rst
+++ b/docs/_templates/autosummary/flax_module.rst
@@ -7,8 +7,8 @@
 
    {% block methods %}
 
-   .. automethod:: __call__ 
-   
+   .. automethod:: __call__
+
    {% if methods %}
    .. rubric:: Methods
 
diff --git a/docs/advanced_topics/contributing.md b/docs/advanced_topics/contributing.md
index d357f9334..1c825c1ec 100644
--- a/docs/advanced_topics/contributing.md
+++ b/docs/advanced_topics/contributing.md
@@ -1,6 +1,6 @@
 # How to Contribute
 
-Everyone can contribute to Flax, and we value everyone's contributions. 
+Everyone can contribute to Flax, and we value everyone's contributions.
 You can contribute in many more ways than just writing code. Answering questions
 on our [Discussions page](https://github.com/google/flax/discussions), helping
 each other, and improving our documentation are extremely valuable to our
@@ -32,7 +32,7 @@ Follow these steps to contribute code:
 
 2. Install Python >=3.6 and `svn` for running the tests (see below).
 
-3. (Optional) Create a virtual environment or a Docker container. See 
+3. (Optional) Create a virtual environment or a Docker container. See
    [`dev/README.md`](https://github.com/google/flax/blob/main/dev/README.md)
    for details on how to setup a Docker Container. To setup a virtual environment,
    run the following:
@@ -41,7 +41,7 @@ Follow these steps to contribute code:
    python3.6 -m virtualenv env
    . env/bin/activate
    ```
-  
+
    This ensures all your dependencies are installed in this environment.
 
 4. `pip install` your fork from source. This allows you to modify the code
@@ -57,7 +57,7 @@ Follow these steps to contribute code:
 
 5. Setup pre-commit hooks, this will run some automated checks during each `git` commit and
    possibly update some files that require changes.
-   
+
    ```bash
    pip install pre-commit
    pre-commit install
@@ -87,7 +87,7 @@ Follow these steps to contribute code:
    ./tests/run_all_tests.sh
    ```
 
-9.  Once your change is done, create a commit as follows 
+9.  Once your change is done, create a commit as follows
    ([how to write a commit message](https://chris.beams.io/posts/git-commit/)):
 
    ```bash
@@ -101,17 +101,17 @@ Follow these steps to contribute code:
    git rebase upstream/main
    ```
 
-11. Finally push your commit on your development branch and create a remote 
+11. Finally push your commit on your development branch and create a remote
    branch in your fork that you can use to create a Pull Request from:
 
    ```bash
    git push --set-upstream origin name-of-change
    ```
-   
+
    After running the command, you should see a Github link in your terminal output that you can click on to create a Pull Request.
    If you do not see this link in the terminal after doing a `git push`, go to the Github web UI; there should be a button there that lets you turn the commit into a Pull Request yourself.
 
-11. Make sure your PR passes the 
+11. Make sure your PR passes the
    [PR checklist](https://github.com/google/flax/blob/main/.github/pull_request_template.md#checklist).
    If so, create a Pull Request from the Flax repository and send it for review.
    Consult [GitHub Help](https://help.github.com/articles/about-pull-requests/)
diff --git a/docs/advanced_topics/convert_pytorch_to_flax.rst b/docs/advanced_topics/convert_pytorch_to_flax.rst
index 1eb453c32..d7b2f0459 100644
--- a/docs/advanced_topics/convert_pytorch_to_flax.rst
+++ b/docs/advanced_topics/convert_pytorch_to_flax.rst
@@ -24,25 +24,25 @@ and the Flax kernel has shape [inC, outC]. Transposing the kernel will do the tr
 .. testcode::
 
   t_fc = torch.nn.Linear(in_features=3, out_features=4)
-  
+
   kernel = t_fc.weight.detach().cpu().numpy()
   bias = t_fc.bias.detach().cpu().numpy()
-  
+
   # [outC, inC] -> [inC, outC]
   kernel = jnp.transpose(kernel, (1, 0))
-  
+
   key = random.PRNGKey(0)
   x = random.normal(key, (1, 3))
-  
+
   variables = {'params': {'kernel': kernel, 'bias': bias}}
   j_fc = nn.Dense(features=4)
   j_out = j_fc.apply(variables, x)
-  
+
   t_x = torch.from_numpy(np.array(x))
   t_out = t_fc(t_x)
   t_out = t_out.detach().cpu().numpy()
-  
-  np.testing.assert_almost_equal(j_out, t_out)  
+
+  np.testing.assert_almost_equal(j_out, t_out)
 
 
 Convolutions
@@ -55,30 +55,30 @@ and the Flax kernel has shape [kH, kW, inC, outC]. Transposing the kernel will d
 .. testcode::
 
   t_conv = torch.nn.Conv2d(in_channels=3, out_channels=4, kernel_size=2, padding='valid')
-  
+
   kernel = t_conv.weight.detach().cpu().numpy()
   bias = t_conv.bias.detach().cpu().numpy()
-  
+
   # [outC, inC, kH, kW] -> [kH, kW, inC, outC]
   kernel = jnp.transpose(kernel, (2, 3, 1, 0))
-  
+
   key = random.PRNGKey(0)
   x = random.normal(key, (1, 6, 6, 3))
-  
+
   variables = {'params': {'kernel': kernel, 'bias': bias}}
   j_conv = nn.Conv(features=4, kernel_size=(2, 2), padding='valid')
   j_out = j_conv.apply(variables, x)
-  
+
   # [N, H, W, C] -> [N, C, H, W]
   t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))
   t_out = t_conv(t_x)
   # [N, C, H, W] -> [N, H, W, C]
   t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))
-  
-  np.testing.assert_almost_equal(j_out, t_out, decimal=6)  
+
+  np.testing.assert_almost_equal(j_out, t_out, decimal=6)
+
 
 
-   
 Convolutions and FC Layers
 --------------------------------
 
@@ -116,7 +116,7 @@ Now, if you want to use the weights from this model in Flax, the corresponding F
 .. testcode::
 
   class JModel(nn.Module):
-      
+
     @nn.compact
     def __call__(self, x):
       x = nn.Conv(features=4, kernel_size=(2, 2), padding='valid', name='conv')(x)
@@ -147,10 +147,10 @@ Other than the transpose operation before reshaping, we can convert the weights
 
   # [outC, inC, kH, kW] -> [kH, kW, inC, outC]
   conv_kernel = jnp.transpose(conv_kernel, (2, 3, 1, 0))
-  
+
   # [outC, inC] -> [inC, outC]
   fc_kernel = jnp.transpose(fc_kernel, (1, 0))
-  
+
   variables = {'params': {'conv': {'kernel': conv_kernel, 'bias': conv_bias},
                           'fc': {'kernel': fc_kernel, 'bias': fc_bias}}}
 
@@ -158,12 +158,12 @@ Other than the transpose operation before reshaping, we can convert the weights
   x = random.normal(key, (1, 6, 6, 3))
 
   j_out = j_model.apply(variables, x)
-  
+
   # [N, H, W, C] -> [N, C, H, W]
   t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))
   t_out = t_model(t_x)
   t_out = t_out.detach().cpu().numpy()
-  
+
   np.testing.assert_almost_equal(j_out, t_out, decimal=6)
 
 
@@ -183,29 +183,29 @@ while Flax multiplies the estimated statistic with ``momentum`` and the new obse
 
   t_bn = torch.nn.BatchNorm2d(num_features=3, momentum=0.1)
   t_bn.eval()
-  
+
   scale = t_bn.weight.detach().cpu().numpy()
   bias = t_bn.bias.detach().cpu().numpy()
   mean = t_bn.running_mean.detach().cpu().numpy()
   var = t_bn.running_var.detach().cpu().numpy()
-  
+
   variables = {'params': {'scale': scale, 'bias': bias},
                'batch_stats': {'mean': mean, 'var': var}}
-  
+
   key = random.PRNGKey(0)
   x = random.normal(key, (1, 6, 6, 3))
-  
+
   j_bn = nn.BatchNorm(momentum=0.9, use_running_average=True)
-  
+
   j_out = j_bn.apply(variables, x)
-  
+
   # [N, H, W, C] -> [N, C, H, W]
   t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))
   t_out = t_bn(t_x)
   # [N, C, H, W] -> [N, H, W, C]
   t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))
-  
-  np.testing.assert_almost_equal(j_out, t_out)  
+
+  np.testing.assert_almost_equal(j_out, t_out)
 
 
 
@@ -226,7 +226,7 @@ operation. ``nn.pool()`` is the core function behind |nn.avg_pool()|_ and |nn.ma
 
 
 .. testcode::
-  
+
   def avg_pool(inputs, window_shape, strides=None, padding='VALID'):
     """
     Pools the input by taking the average over a window.
@@ -234,25 +234,25 @@ operation. ``nn.pool()`` is the core function behind |nn.avg_pool()|_ and |nn.ma
     consider the padded zero's for the average computation.
     """
     assert len(window_shape) == 2
-  
+
     y = nn.pool(inputs, 0., jax.lax.add, window_shape, strides, padding)
     counts = nn.pool(jnp.ones_like(inputs), 0., jax.lax.add, window_shape, strides, padding)
-    y = y / counts 
+    y = y / counts
     return y
 
 
   key = random.PRNGKey(0)
   x = random.normal(key, (1, 6, 6, 3))
-  
+
   j_out = avg_pool(x, window_shape=(2, 2), strides=(1, 1), padding=((1, 1), (1, 1)))
   t_pool = torch.nn.AvgPool2d(kernel_size=2, stride=1, padding=1, count_include_pad=False)
-  
+
   # [N, H, W, C] -> [N, C, H, W]
   t_x = torch.from_numpy(np.transpose(np.array(x), (0, 3, 1, 2)))
   t_out = t_pool(t_x)
   # [N, C, H, W] -> [N, H, W, C]
   t_out = np.transpose(t_out.detach().cpu().numpy(), (0, 2, 3, 1))
-  
+
   np.testing.assert_almost_equal(j_out, t_out)
 
 
diff --git a/docs/advanced_topics/lift.md b/docs/advanced_topics/lift.md
index c14f62bbd..47f92f739 100644
--- a/docs/advanced_topics/lift.md
+++ b/docs/advanced_topics/lift.md
@@ -147,7 +147,7 @@ For example, we want to share all parameter variables when we add a batch dimens
 At the same time we can write generic code that uses transformations without knowing exactly what kind of variables the submodules will create.
 Collections thus strike a balance between fine-grained control and generality.
 We also avoid brittle string matching code that loops over all variables and tries to split up collections in an ad-hoc way based on
-naming conventions like: target all variables with the name prefix "kernel". 
+naming conventions like: target all variables with the name prefix "kernel".
 If more fine-grained control is necessary a user can simply split up a set of variables over multiple collections that should be handled differently.
 
 
@@ -160,7 +160,7 @@ The reason for this choice is that there are many different transformations with
 For example: `vmap` has broadcasted and vectorized arguments, while `scan` has scan, carry, and broadcast arguments.
 A variable would have to define its behavior for all these transformations otherwise a `Module` would not be compatible with
 these transformations. Alternatively, we would have to make default decisions for how transformations are handled.
-However, this could lead to silent bugs because the behavior might not actually be valid given the users intent. 
+However, this could lead to silent bugs because the behavior might not actually be valid given the users intent.
 
 The lift package also provides a general purpose `transform`, which allows an arbitrary function to transform a variable collection.
 For example, this can be used to tie the weights in a tied auto-encoder by transposing the weights.
@@ -231,7 +231,7 @@ The `pack(fn, in_vars, out_vars, rngs)` API goes through the following stages:
     After we have used the lifted scopes we have to retrieve the updated variables (PRNG sequences can simply be discarded).
     pack passes the `repack_fn` to support this.
     This stage is similar to stage 2 except that we only lift variables and immutable variables are ignored.
-    Immutable variables cannot be updated. Therefore, they should not be returned from the transformed function. 
+    Immutable variables cannot be updated. Therefore, they should not be returned from the transformed function.
 
 6. *Commit stage*
 
diff --git a/docs/advanced_topics/linen_upgrade_guide.rst b/docs/advanced_topics/linen_upgrade_guide.rst
index 5e1447a59..8c8258933 100644
--- a/docs/advanced_topics/linen_upgrade_guide.rst
+++ b/docs/advanced_topics/linen_upgrade_guide.rst
@@ -28,11 +28,11 @@ can use this upgrade guide to upgrade it to Linen.
 Defining Simple Modules
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
-  
+
   from flax import nn
 
   class Dense(base.Module):
@@ -50,8 +50,8 @@ Defining Simple Modules
         bias = self.param(
           'bias', (features,), bias_init)
         y = y + bias
-      return y  
-  
+      return y
+
     return new_state, metrics
   ---
   from flax import linen as nn  # [1] #!
@@ -94,11 +94,11 @@ Defining Simple Modules
 Using Modules inside other Modules
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
-  
+
   class Encoder(nn.Module):
 
     def apply(self, x):
@@ -126,11 +126,11 @@ Using Modules inside other Modules
 Sharing submodules and defining multiple methods
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
-  :sync: 
-  
+  :sync:
+
   class AutoEncoder(nn.Module):
     def _create_submodules(self):
       return Decoder.shared(name="encoder")
@@ -178,13 +178,13 @@ Sharing submodules and defining multiple methods
 ``Module.partial`` inside other modules
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
-  
+
   # no import #!
-   
+
   class ResNet(nn.Module):
     """ResNetV1."""
 
@@ -204,7 +204,7 @@ Sharing submodules and defining multiple methods
       x = norm(x, name='bn_init')
 
       # [...]
-      return x 
+      return x
   ---
   from functools import partial  #!
 
@@ -230,12 +230,12 @@ Sharing submodules and defining multiple methods
       return x
 
 Use normal ``functools.partial`` instead of ``Module.partial``. The rest stays
-the same. 
+the same.
 
 Top-level training code patterns
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
@@ -289,7 +289,7 @@ Top-level training code patterns
 1. We no longer use the ``Model`` abstraction -- instead we pass parameters
    around directly, usually encapsulated in a `Train State`_ object, which can
    directly be passed to JAX transformations.
-   
+
 2. To compute initial parameters, construct a module instance and call |init|_
    or |init_with_output|_. We haven't ported over ``init_by_shape`` because this
    function did some magic we did not like (it evaluated the function by shape.
@@ -313,11 +313,11 @@ Top-level training code patterns
 Non-trainable variables ("state"): Use within Modules
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
-  
+
   class BatchNorm(nn.Module):
     def apply(self, x, ...):
       # [...]
@@ -325,7 +325,7 @@ Non-trainable variables ("state"): Use within Modules
         'mean', (x.shape[-1], ), initializers.zeros)
       ra_var = self.state(
         'var', (x.shape[-1], ), initializers.ones)
-      # [...] 
+      # [...]
   ---
   class BatchNorm(nn.Module):
     def __call__(self, x):
@@ -345,7 +345,7 @@ using JAX transformations inside modules.
 Non-trainable variables ("state"): Top-level training code patterns
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
@@ -386,7 +386,7 @@ Non-trainable variables ("state"): Top-level training code patterns
     variables = {'params': params, 'batch_stats': batch_stats}  # [2] #!
     logits, new_variables = ResNet(train=true).apply(
       variables, batch['image'], mutable=['batch_stats'])  # [3] #!
-    new_batch_stats = new_variables['batch_stats'] 
+    new_batch_stats = new_variables['batch_stats']
     # [...]
 
 
@@ -396,7 +396,7 @@ Non-trainable variables ("state"): Top-level training code patterns
     logits = ResNet(train=False).apply(
       variables, batch['image'], mutable=False)  # [4] #!
     return compute_metrics(logits, batch['label'])
-  
+
 1. |init|_ returns a variable dict, e.g. ``{"param": ..., "batch_stats": ...}``
    (see `Variable documentation`_).
 
@@ -432,11 +432,11 @@ TODO: Add an example here how to load a new ``TrainState`` object.
 Randomness
 --------------------------------
 
-.. codediff:: 
+.. codediff::
   :title_left: Old Flax
   :title_right: Linen
   :sync:
-  
+
   def dropout(inputs, rate, deterministic=False):
     keep_prob = 1. - rate
     if deterministic:
diff --git a/docs/advanced_topics/module_lifecycle.rst b/docs/advanced_topics/module_lifecycle.rst
index f884f5948..76bc42f63 100644
--- a/docs/advanced_topics/module_lifecycle.rst
+++ b/docs/advanced_topics/module_lifecycle.rst
@@ -36,7 +36,7 @@ Let's start with a high-level overview of the Module lifecycle. First, define a
       self.hidden = nn.Dense(self.hidden_size)
       self.out = nn.Dense(self.out_size)
 
-    # 3. User methods  
+    # 3. User methods
     def __call__(self, x):
       a = self.hidden(x)
       h = nn.relu(a)
@@ -378,4 +378,4 @@ Imagine we want to call just the decoder using `auto_encoder.decoder.apply(decod
 
 There are two possible solutions to make this use case more ergonomic. First, setup could be made to run immediately after construction before it becomes bound. This means you can still create sub modules but you can no longer define or manipulate variables. Therefore, this would be a breaking change and it would require a new API for defining variables lazily
 
-Alternatively, an additional special method could be introduced that runs right away after Module construction and before it becomes bound. In this case, the ``setup`` method would preserve its original semantics. 
+Alternatively, an additional special method could be introduced that runs right away after Module construction and before it becomes bound. In this case, the ``setup`` method would preserve its original semantics.
diff --git a/docs/advanced_topics/optax_update_guide.rst b/docs/advanced_topics/optax_update_guide.rst
index e3792344b..8da1a4bc5 100644
--- a/docs/advanced_topics/optax_update_guide.rst
+++ b/docs/advanced_topics/optax_update_guide.rst
@@ -22,7 +22,7 @@ https://optax.readthedocs.io/en/latest/optax-101.html
   import optax
 
   # Note: this is the minimal code required to make below code run. See in the
-  # Colab linked above for a more meaningful definition of datasets etc. 
+  # Colab linked above for a more meaningful definition of datasets etc.
   batch = {'image': jnp.ones([1, 28, 28, 1]), 'label': jnp.array([0])}
   ds_train = [batch]
   get_ds_train = lambda: [batch]
@@ -106,7 +106,7 @@ generic building blocks.
   #
 
   tx = optax.chain(
-      # 1. Step: keep a trace of past updates and add to gradients. 
+      # 1. Step: keep a trace of past updates and add to gradients.
       optax.trace(decay=momentum),
       # 2. Step: multiply result from step 1 with negative learning rate.
       # Note that `optax.apply_updates()` simply adds the final updates to the
diff --git a/docs/conf.py b/docs/conf.py
index 0c1fd9b06..506f83c68 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -128,7 +128,7 @@
 # files that will not be executed.
 myst_enable_extensions = ['dollarmath']
 nb_execution_excludepatterns = [
-  'getting_started.ipynb', # <-- times out 
+  'getting_started.ipynb', # <-- times out
   'optax_update_guide.ipynb', # <-- requires flax<=0.5.3
 ]
 # raise exceptions on execution so CI can catch errors
diff --git a/docs/examples.rst b/docs/examples.rst
index 39ee02ddf..2d4350f15 100644
--- a/docs/examples.rst
+++ b/docs/examples.rst
@@ -6,16 +6,16 @@ Core examples
 ##############
 
 
-Core examples are hosted on the Flax repo in the `examples <https://github.com/google/flax/tree/main/examples>`__ 
+Core examples are hosted on the Flax repo in the `examples <https://github.com/google/flax/tree/main/examples>`__
 directory.
 
 Each example is designed to be **self-contained and easily forkable**, while
 reproducing relevant results in different areas of machine learning.
 
-As discussed in `#231 <https://github.com/google/flax/issues/231>`__, we decided 
-to go for a standard pattern for all examples including the simplest ones (like MNIST). 
-This makes every example a bit more verbose, but once you know one example, you 
-know the structure of all of them. Having unit tests and integration tests is also 
+As discussed in `#231 <https://github.com/google/flax/issues/231>`__, we decided
+to go for a standard pattern for all examples including the simplest ones (like MNIST).
+This makes every example a bit more verbose, but once you know one example, you
+know the structure of all of them. Having unit tests and integration tests is also
 very useful when you fork these examples.
 
 Some of the examples below have a link "Interactive🕹" that lets you run them
@@ -105,7 +105,7 @@ of examples, in many cases with pre-trained weights:
   `Flax/JAX models <https://huggingface.co/models?library=jax&sort=downloads>`__ in
   their repository.
 
-- `🥑 DALLE Mini <https://huggingface.co/dalle-mini>`__ is a Transformer-based 
+- `🥑 DALLE Mini <https://huggingface.co/dalle-mini>`__ is a Transformer-based
   text-to-image model implemented in JAX/Flax that follows the ideas from the
   original `DALLE <https://openai.com/blog/dall-e/>`__ paper by OpenAI.
 
@@ -203,20 +203,20 @@ Tutorials
 Contributing Policy
 ********************
 
-If you are interested in adding a project to the Community Examples section, take the following 
+If you are interested in adding a project to the Community Examples section, take the following
 into consideration:
 
-* **Examples**: examples should contain a README that is helpful, clear, and makes it easy to run 
+* **Examples**: examples should contain a README that is helpful, clear, and makes it easy to run
   the code. The code itself should be easy to follow.
-* **Tutorials**: tutorials must preferably be runnable notebooks, be well written, and discuss 
+* **Tutorials**: tutorials must preferably be runnable notebooks, be well written, and discuss
   an interesting topic. Also, the tutorial's content must be different from the existing
   guides in the Flax documentation and other community examples to be considered for inclusion.
 * **Models**: repositories with models ported to Flax must provide at least one of the following:
-  
-  * Metrics that are comparable to the original work when the model is trained to completion. Having 
+
+  * Metrics that are comparable to the original work when the model is trained to completion. Having
     available plots of the metric's history during training is highly encouraged.
-  * Tests to verify numerical equivalence against a well known implementation (same inputs 
+  * Tests to verify numerical equivalence against a well known implementation (same inputs
     + weights = same outputs) preferably using pretrained weights.
 
-On all cases above, code should work with the latest stable version of packages like ``jax``, 
+On all cases above, code should work with the latest stable version of packages like ``jax``,
 ``flax``, and ``optax``, and make substantial use of Flax.
diff --git a/docs/flip/1777-default-dtype.md b/docs/flip/1777-default-dtype.md
index f0949c328..6344b6bb0 100644
--- a/docs/flip/1777-default-dtype.md
+++ b/docs/flip/1777-default-dtype.md
@@ -60,7 +60,7 @@ class Dense(nn.Module):
 
  @nn.compact
  def __call__(self, x):
-   kernel = self.param("kernel", 
+   kernel = self.param("kernel",
                        self.kernel_init,
                        (x.shape[-1], self.features), self.param_dtype)
    bias = self.param("bias", self.bias_init, (self.features,), self.param_dtype)
diff --git a/docs/flip/2434-general-metadata.md b/docs/flip/2434-general-metadata.md
index d432eaa31..6cb67023f 100644
--- a/docs/flip/2434-general-metadata.md
+++ b/docs/flip/2434-general-metadata.md
@@ -9,7 +9,7 @@
 
 ## Summary
 
-This FLIP proposes to extend Flax's variable collections with a generic axis metadata API. 
+This FLIP proposes to extend Flax's variable collections with a generic axis metadata API.
 The core of the API is an abstract base class that is recognized by lifting transformations that can add an axis (vmap, scan).
 Users can extend the base class to keep track of per-axis metadata in a way that works with lifted transformations.
 
@@ -153,7 +153,7 @@ To illustrate this, let's consider a metadata class that keeps track of Partitio
 class Partitioned(flax.struct.PyTreeNode, AxisMetadata):
   value: Any
   names: Tuple[Optional[str], ...] = flax.struct.field(pytree_node=False)
-  
+
   def add_axis(self, index: int, params: Dict[Any, Any]) -> TAxisMetadata:
     axis_name = self._get_partition_name(params)
     names = list(self.names)
@@ -172,7 +172,7 @@ def with_partitioning(init_fn, names):
   return wrapper
 ```
 
-Here we also defined a small utility called ``with_partitioning`` that we can use to wrap existing initialzers to add metadata: 
+Here we also defined a small utility called ``with_partitioning`` that we can use to wrap existing initialzers to add metadata:
 
 
 ```python
@@ -198,7 +198,7 @@ def to_sharding_spec(x):
     # fully replicated
     return PartitionSpec()
 
-# Result: {"params": {"kernel": PartitionSpec(None, "data"), bias: PartitionSpec()}} 
+# Result: {"params": {"kernel": PartitionSpec(None, "data"), bias: PartitionSpec()}}
 variables_pspec = jax.tree_map(to_sharding_spec, variables, is_leaf=lambda x: isinstance(x, Partitioned))
 ```
 
diff --git a/docs/getting_started.ipynb b/docs/getting_started.ipynb
index 6d801c461..e053e5e88 100644
--- a/docs/getting_started.ipynb
+++ b/docs/getting_started.ipynb
@@ -601,7 +601,7 @@
     "  rng, input_rng = jax.random.split(rng)\n",
     "  # Run an optimization step over a training batch\n",
     "  state = train_epoch(state, train_ds, batch_size, epoch, input_rng)\n",
-    "  # Evaluate on the test set after each training epoch \n",
+    "  # Evaluate on the test set after each training epoch\n",
     "  test_loss, test_accuracy = eval_model(state.params, test_ds)\n",
     "  print(' test epoch: %d, loss: %.2f, accuracy: %.2f' % (\n",
     "      epoch, test_loss, test_accuracy * 100))"
diff --git a/docs/getting_started.md b/docs/getting_started.md
index db14969f1..fccb06fd9 100644
--- a/docs/getting_started.md
+++ b/docs/getting_started.md
@@ -391,7 +391,7 @@ for epoch in range(1, num_epochs + 1):
   rng, input_rng = jax.random.split(rng)
   # Run an optimization step over a training batch
   state = train_epoch(state, train_ds, batch_size, epoch, input_rng)
-  # Evaluate on the test set after each training epoch 
+  # Evaluate on the test set after each training epoch
   test_loss, test_accuracy = eval_model(state.params, test_ds)
   print(' test epoch: %d, loss: %.2f, accuracy: %.2f' % (
       epoch, test_loss, test_accuracy * 100))
diff --git a/docs/guides/extracting_intermediates.rst b/docs/guides/extracting_intermediates.rst
index d4e54fdb2..ef7295c1a 100644
--- a/docs/guides/extracting_intermediates.rst
+++ b/docs/guides/extracting_intermediates.rst
@@ -43,10 +43,10 @@ Store intermediate values in a new variable collection
 The CNN can be augmented with calls to ``sow`` to store intermediates as following:
 
 
-.. codediff:: 
+.. codediff::
   :title_left: Default CNN
   :title_right: CNN using sow API
-  
+
   class CNN(nn.Module):
     @nn.compact
     def __call__(self, x):
@@ -223,7 +223,7 @@ Use ``Sequential``
 
 You could also define ``CNN`` using a simple implementation of a ``Sequential`` combinator (this is quite common in more stateful approaches). This may be useful
 for very simple models and gives you arbitrary model
-surgery. But it can be very limiting -- if you even want to add one conditional, you are 
+surgery. But it can be very limiting -- if you even want to add one conditional, you are
 forced to refactor away from ``Sequential`` and structure
 your model more explicitly.
 
diff --git a/docs/guides/flax_basics.ipynb b/docs/guides/flax_basics.ipynb
index 5e3422a94..4806d938c 100644
--- a/docs/guides/flax_basics.ipynb
+++ b/docs/guides/flax_basics.ipynb
@@ -56,7 +56,7 @@
    ],
    "source": [
     "# Install the latest JAXlib version.\n",
-    "!pip install --upgrade -q pip jax jaxlib \n",
+    "!pip install --upgrade -q pip jax jaxlib\n",
     "# Install Flax at head:\n",
     "!pip install --upgrade -q git+https://github.com/google/flax.git"
    ]
@@ -869,7 +869,7 @@
     "id": "MKyhfzVpzC94"
    },
    "source": [
-    "Here, we see how to both declare and assign a parameter to the model using the `self.param` method. It takes as input `(name, init_fn, *init_args)` : \n",
+    "Here, we see how to both declare and assign a parameter to the model using the `self.param` method. It takes as input `(name, init_fn, *init_args)` :\n",
     "\n",
     "*   `name` is simply the name of the parameter that will end up in the parameter structure.\n",
     "*   `init_fn` is a function with input `(PRNGKey, *init_args)` returning an Array, with `init_args` being the arguments needed to call the initialisation function.\n",
diff --git a/docs/guides/flax_basics.md b/docs/guides/flax_basics.md
index de83e1e73..af4874d75 100644
--- a/docs/guides/flax_basics.md
+++ b/docs/guides/flax_basics.md
@@ -41,7 +41,7 @@ outputId: e30aa464-fa52-4f35-df96-716c68a4b3ee
 tags: [skip-execution]
 ---
 # Install the latest JAXlib version.
-!pip install --upgrade -q pip jax jaxlib 
+!pip install --upgrade -q pip jax jaxlib
 # Install Flax at head:
 !pip install --upgrade -q git+https://github.com/google/flax.git
 ```
@@ -469,7 +469,7 @@ print('output:\n', y)
 
 +++ {"id": "MKyhfzVpzC94"}
 
-Here, we see how to both declare and assign a parameter to the model using the `self.param` method. It takes as input `(name, init_fn, *init_args)` : 
+Here, we see how to both declare and assign a parameter to the model using the `self.param` method. It takes as input `(name, init_fn, *init_args)` :
 
 *   `name` is simply the name of the parameter that will end up in the parameter structure.
 *   `init_fn` is a function with input `(PRNGKey, *init_args)` returning an Array, with `init_args` being the arguments needed to call the initialisation function.
diff --git a/docs/guides/index.rst b/docs/guides/index.rst
index 5a44d5c16..9ee947ee5 100644
--- a/docs/guides/index.rst
+++ b/docs/guides/index.rst
@@ -18,4 +18,3 @@ Guides
 
    ensembling
    full_eval
-   
\ No newline at end of file
diff --git a/docs/guides/jax_for_the_impatient.ipynb b/docs/guides/jax_for_the_impatient.ipynb
index b10909013..edbc51410 100644
--- a/docs/guides/jax_for_the_impatient.ipynb
+++ b/docs/guides/jax_for_the_impatient.ipynb
@@ -37,7 +37,7 @@
     "import jax\n",
     "from jax import numpy as jnp, random\n",
     "\n",
-    "import numpy as np # We import the standard NumPy library "
+    "import numpy as np # We import the standard NumPy library"
    ]
   },
   {
@@ -46,7 +46,7 @@
     "id": "Z5BLL6v_JUSI"
    },
    "source": [
-    "`jax.numpy` is the NumPy-like API that needs to be imported, and we will also use `jax.random` to generate some data to work on. \n",
+    "`jax.numpy` is the NumPy-like API that needs to be imported, and we will also use `jax.random` to generate some data to work on.\n",
     "\n",
     "Let's start by generating some matrices, and then try matrix multiplication."
    ]
diff --git a/docs/guides/jax_for_the_impatient.md b/docs/guides/jax_for_the_impatient.md
index 9beb486fb..ebc6d88f0 100644
--- a/docs/guides/jax_for_the_impatient.md
+++ b/docs/guides/jax_for_the_impatient.md
@@ -34,12 +34,12 @@ Let's start by exploring the NumPy API coming from JAX and the main differences
 import jax
 from jax import numpy as jnp, random
 
-import numpy as np # We import the standard NumPy library 
+import numpy as np # We import the standard NumPy library
 ```
 
 +++ {"id": "Z5BLL6v_JUSI"}
 
-`jax.numpy` is the NumPy-like API that needs to be imported, and we will also use `jax.random` to generate some data to work on. 
+`jax.numpy` is the NumPy-like API that needs to be imported, and we will also use `jax.random` to generate some data to work on.
 
 Let's start by generating some matrices, and then try matrix multiplication.
 
diff --git a/docs/guides/lr_schedule.rst b/docs/guides/lr_schedule.rst
index 6c87b064f..71db202d6 100644
--- a/docs/guides/lr_schedule.rst
+++ b/docs/guides/lr_schedule.rst
@@ -29,10 +29,10 @@ We will show you how to...
   import ml_collections
   from absl import logging
 
-   
+
   class CNN(nn.Module):
     """A simple CNN model."""
-   
+
     @nn.compact
     def __call__(self, x):
       x = nn.Conv(features=32, kernel_size=(3, 3))(x)
@@ -46,16 +46,16 @@ We will show you how to...
       x = nn.relu(x)
       x = nn.Dense(features=10)(x)
       return x
-   
+
   def get_dummy_data(ds_size):
     image = np.random.rand(ds_size, 28, 28, 1)
     label = np.random.randint(low=0, high=10, size=(ds_size,))
     return {'image': image, 'label': label}
-  
+
   def get_config():
     """Get the default hyperparameter configuration."""
     config = ml_collections.ConfigDict()
-  
+
     config.learning_rate = 0.001
     config.momentum = 0.9
     config.batch_size = 128
@@ -76,7 +76,7 @@ We will show you how to...
 
 
 .. testcode::
-  
+
   def create_learning_rate_fn(config, base_learning_rate, steps_per_epoch):
     """Creates learning rate schedule."""
     warmup_fn = optax.linear_schedule(
@@ -98,11 +98,11 @@ For example using this schedule on MNIST would require changing the ``train_step
 .. |Optax| replace:: ``Optax``
 .. _Optax: https://optax.readthedocs.io/en/latest/api.html#optimizer-schedules
 
-.. codediff:: 
+.. codediff::
   :title_left: Default learning rate
   :title_right: Learning rate schedule
   :sync:
-  
+
   @jax.jit
   def train_step(state, batch):
     def loss_fn(params):
@@ -114,8 +114,8 @@ For example using this schedule on MNIST would require changing the ``train_step
     (_, logits), grads = grad_fn(state.params)
     new_state = state.apply_gradients(grads=grads)
     metrics = compute_metrics(logits, batch['label'])
-  
-  
+
+
     return new_state, metrics
   ---
   @functools.partial(jax.jit, static_argnums=2) #!
@@ -139,7 +139,7 @@ And the ``train_epoch`` function:
   :title_left: Default learning rate
   :title_right: Learning rate schedule
   :sync:
-  
+
   def train_epoch(state, train_ds, batch_size, epoch, rng):
     """Trains for a single epoch."""
     train_ds_size = len(train_ds['image'])
@@ -152,16 +152,16 @@ And the ``train_epoch`` function:
       batch = {k: v[perm, ...] for k, v in train_ds.items()}
       state, metrics = train_step(state, batch)
       batch_metrics.append(metrics)
-  
+
     # compute mean of metrics across each batch in epoch.
     batch_metrics = jax.device_get(batch_metrics)
     epoch_metrics = {
         k: np.mean([metrics[k] for metrics in batch_metrics])
         for k in batch_metrics[0]}
-  
+
     logging.info('train epoch: %d, loss: %.4f, accuracy: %.2f', epoch,
                  epoch_metrics['loss'], epoch_metrics['accuracy'] * 100)
-  
+
     return state, epoch_metrics
   ---
   def train_epoch(state, train_ds, batch_size, epoch, learning_rate_fn, rng): #!
@@ -176,16 +176,16 @@ And the ``train_epoch`` function:
       batch = {k: v[perm, ...] for k, v in train_ds.items()}
       state, metrics = train_step(state, batch, learning_rate_fn) #!
       batch_metrics.append(metrics)
-  
+
     # compute mean of metrics across each batch in epoch.
     batch_metrics = jax.device_get(batch_metrics)
     epoch_metrics = {
         k: np.mean([metrics[k] for metrics in batch_metrics])
         for k in batch_metrics[0]}
-  
+
     logging.info('train epoch: %d, loss: %.4f, accuracy: %.2f', epoch,
                  epoch_metrics['loss'], epoch_metrics['accuracy'] * 100)
-  
+
     return state, epoch_metrics
 
 
@@ -217,11 +217,11 @@ And the ``create_train_state`` function:
 .. testcleanup::
 
   config = get_config()
-  
+
   train_ds_size = config.train_ds_size
   steps_per_epoch = train_ds_size // config.batch_size
   learning_rate_fn = create_learning_rate_fn(config, config.learning_rate, steps_per_epoch)
-  
+
   rng = jax.random.PRNGKey(0)
   state = create_train_state(rng, config, learning_rate_fn)
 
@@ -231,5 +231,5 @@ And the ``create_train_state`` function:
 
   assert 'accuracy' in epoch_metrics and 'learning_rate' in epoch_metrics
 
-   
+
 
diff --git a/docs/guides/model_surgery.rst b/docs/guides/model_surgery.rst
index 363ab2376..882fcfb56 100644
--- a/docs/guides/model_surgery.rst
+++ b/docs/guides/model_surgery.rst
@@ -13,7 +13,7 @@ Model Surgery
   import jax
   import optax
 
-We will show how to get a flat dict of all the tensors, and then go back to a 
+We will show how to get a flat dict of all the tensors, and then go back to a
 nested, frozen dict. This will be demonstrated for both Flax modules and optimizers.
 
 Surgery with Flax Modules
@@ -154,7 +154,7 @@ optimizer state that mirrors the original state.
       Dense_0: { bias: (256,), kernel: (3136, 256), },
       Dense_1: { bias: (10,), kernel: (256, 10), },
   })), EmptyState())
-  
+
 The pytrees inside the optimizer state follow the same structure as the
 parameters and can be flattened / modified exactly the same way
 
@@ -167,7 +167,7 @@ parameters and can be flattened / modified exactly the same way
 
 .. testoutput::
   :options: +NORMALIZE_WHITESPACE
-  
+
   {('Conv_0', 'bias'): (32,),
    ('Conv_0', 'kernel'): (3, 3, 1, 32),
    ('Conv_1', 'bias'): (64,),
diff --git a/docs/guides/setup_or_nncompact.rst b/docs/guides/setup_or_nncompact.rst
index 2ead5b0ee..1577274c6 100644
--- a/docs/guides/setup_or_nncompact.rst
+++ b/docs/guides/setup_or_nncompact.rst
@@ -8,34 +8,34 @@ can be defined in two ways:
 
    Assign submodules or variables to ``self.<attr>`` inside a
    :meth:`setup <flax.linen.Module.setup>` method. Then use the submodules
-   and variables assigned to ``self.<attr>`` in ``setup`` from 
+   and variables assigned to ``self.<attr>`` in ``setup`` from
    any "forward pass" method defined on the class.
    This resembles how modules are defined in PyTorch.
 
 2. **In-line** (using ``nn.compact``):
-  
+
    Write your network's logic directly within a single "forward pass" method annotated
    with :meth:`nn.compact <flax.linen.compact>`. This allows you to define your whole module
-   in a single method, and "co-locate" submodules and variables next to 
+   in a single method, and "co-locate" submodules and variables next to
    where they are used.
 
 **Both of these approaches are perfectly valid, behave the same way, and interoperate with all of Flax**.
 
 Here is a short example of a module defined in both ways, with exactly
-the same functionality. 
+the same functionality.
 
 .. testsetup::
 
   import flax.linen as nn
 
-.. codediff:: 
+.. codediff::
   :title_left: Using ``setup``
   :title_right: Using ``nn.compact``
-  
+
   class MLP(nn.Module):
     def setup(self):
       # Submodule names are derived by the attributes you assign to. In this
-      # case, "dense1" and "dense2". This follows the logic in PyTorch. 
+      # case, "dense1" and "dense2". This follows the logic in PyTorch.
       self.dense1 = nn.Dense(32)
       self.dense2 = nn.Dense(32)
 
diff --git a/docs/index.rst b/docs/index.rst
index 832033575..a9d43f70d 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -35,11 +35,11 @@ Features
          :class-title: sd-fs-5
 
          .. div:: sd-font-normal
-            
-            Flax is designed for correctness and safety. Thanks to its immutable Modules 
-            and Functional API, Flax helps mitigate bugs that araise when handling state 
+
+            Flax is designed for correctness and safety. Thanks to its immutable Modules
+            and Functional API, Flax helps mitigate bugs that araise when handling state
             in JAX.
-   
+
    .. grid-item::
       :columns: 12 12 12 6
 
@@ -49,10 +49,10 @@ Features
          :class-title: sd-fs-5
 
          .. div:: sd-font-normal
-            
+
             Flax grants more fine grained control and expressivity than most Neural Network
             frameworks via its Variable Collections, RNG Collections and Mutability conditions.
-   
+
    .. grid-item::
       :columns: 12 12 12 6
 
@@ -62,9 +62,9 @@ Features
          :class-title: sd-fs-5
 
          .. div:: sd-font-normal
-            
+
             Flax's functional API radically redefines what Modules can do via lifted transformations like vmap, scan, etc, while also enabling seamless integration with other JAX libraries like Optax and Chex.
-   
+
    .. grid-item::
       :columns: 12 12 12 6
 
@@ -74,7 +74,7 @@ Features
          :class-title: sd-fs-5
 
          .. div:: sd-font-normal
-            
+
             Flax's :meth:`compact <flax.linen.compact>` Modules enables submodules to be defined directly at their callsite, leading to code that is easier to read and avoids repetition.
 
 
@@ -125,8 +125,8 @@ Learn more
    .. grid-item::
       :columns: 6 6 6 4
 
-      
-      .. card:: :material-regular:`rocket_launch;2em` Getting Started 
+
+      .. card:: :material-regular:`rocket_launch;2em` Getting Started
          :class-card: sd-text-black sd-bg-light
          :link: getting_started.html
 
@@ -179,83 +179,83 @@ Notable examples include:
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             NLP and Computer Vision models
-   
+
    .. grid-item::
       :columns: 6 6 6 4
-      
+
       .. card:: `🥑 DALLE Mini <https://huggingface.co/dalle-mini>`__
          :class-card: sd-text-black sd-border-0
          :shadow: none
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             Model for Text-to-Image generation
 
    .. grid-item::
       :columns: 6 6 6 4
-      
+
       .. card:: `PaLM <https://ai.googleblog.com/2022/04/pathways-language-model-palm-scaling-to.html>`__
          :class-card: sd-text-black sd-border-0
          :shadow: none
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             540 Billion parameter model for text generation
 
    .. grid-item::
       :columns: 6 6 6 4
-      
+
       .. card:: `Imagen <https://imagen.research.google>`__
          :class-card: sd-text-black sd-border-0
          :shadow: none
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             Text-to-Image Diffusion Models
-   
+
    .. grid-item::
       :columns: 6 6 6 4
-      
+
       .. card:: `Big Vision <https://github.com/google-research/big_vision>`__
          :class-card: sd-text-black sd-border-0
          :shadow: none
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             Large scale Computer Vision models
-   
+
    .. grid-item::
       :columns: 6 6 6 4
-      
+
       .. card:: `T5x <https://github.com/google-research/t5x>`__
          :class-card: sd-text-black sd-border-0
          :shadow: none
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             Large Language Models
-   
+
    .. grid-item::
       :columns: 6 6 6 4
-      
+
       .. card:: `Brax <https://github.com/google/brax>`__
          :class-card: sd-text-black sd-border-0
          :shadow: none
          :class-title: sd-text-center sd-fs-5
 
          .. div:: sd-text-center sd-font-italic
-            
+
             On-device differentiable RL environments
 
 
-      
+
 
 .. toctree::
    :hidden:
diff --git a/docs/mission.md b/docs/mission.md
index aee85ba2d..c6a2b8884 100644
--- a/docs/mission.md
+++ b/docs/mission.md
@@ -24,7 +24,7 @@ JAX is an excellent basis for current research in many fields in numerical compu
 
 We, on the Flax team, aim to ensure the JAX ecosystem has the necessary tools and guidance to facilitate and maximize this new space of research directions in the form of libraries, integrated examples for both educational purposes, as well as ready-to-fork reference implementations, guides, support channels, and best practices.
 
-Flax aims to support researchers' needs, while exposing the full power of JAX, rather than encapsulating it away. We think that this is the best path towards ground-breaking research that builds on JAX's unique capabilities. Some recent examples of innovative research projects on top of Flax include [NetKet](https://www.netket.org/) for many-body quantum systems, [T5X](https://github.com/google-research/t5x) for very large scale language model and [Brax](https://github.com/google/brax) for differentiable on-device rigid body simulations. 
+Flax aims to support researchers' needs, while exposing the full power of JAX, rather than encapsulating it away. We think that this is the best path towards ground-breaking research that builds on JAX's unique capabilities. Some recent examples of innovative research projects on top of Flax include [NetKet](https://www.netket.org/) for many-body quantum systems, [T5X](https://github.com/google-research/t5x) for very large scale language model and [Brax](https://github.com/google/brax) for differentiable on-device rigid body simulations.
 
 It is challenging to design a library that is ergonomic for researchers, while building on a purely functional library like JAX, which may be confusing to some. But we see this challenge as **an opportunity**. How do you design a smooth user experience where the complexity is progressively disclosed? How do you help users avoid "silent footguns", such as accidentally mutating batch norm stats during prediction? How do you make sure errors guide users to the right solution, even when there may be multiple correct ones? How do you design examples and tutorials that are simple enough to be educational, while being good enough to act as forkable starter codebases? These are the types of questions that inform our design process.
 
@@ -42,6 +42,6 @@ Does this mission excite you? If so, take a look at [our job posting](https://ca
 
 With great hopes for the future,
 
-Marc, Jonathan, George, Bertrand, Avital, Anselm, Andreas, 
+Marc, Jonathan, George, Bertrand, Avital, Anselm, Andreas,
 
 on behalf of The Flax Team
diff --git a/docs/notebooks/full_eval.ipynb b/docs/notebooks/full_eval.ipynb
index ff484cf98..014745465 100644
--- a/docs/notebooks/full_eval.ipynb
+++ b/docs/notebooks/full_eval.ipynb
@@ -50,7 +50,7 @@
     }
    ],
    "source": [
-    "!pip install -q chex einops \n",
+    "!pip install -q chex einops\n",
     "# tfds.split_for_jax_process() was added in 4.5.1\n",
     "!pip install -q tensorflow_datasets -U\n",
     "# flax.jax_utils.pad_shard_unpad() is only available at HEAD\n",
diff --git a/docs/notebooks/full_eval.md b/docs/notebooks/full_eval.md
index ac17fbb46..f9cd9f7f4 100644
--- a/docs/notebooks/full_eval.md
+++ b/docs/notebooks/full_eval.md
@@ -33,7 +33,7 @@ id: 62DTHYCYHWp1
 outputId: b38d096f-58db-4d61-effa-eafa4c732826
 tags: [skip-execution]
 ---
-!pip install -q chex einops 
+!pip install -q chex einops
 # tfds.split_for_jax_process() was added in 4.5.1
 !pip install -q tensorflow_datasets -U
 # flax.jax_utils.pad_shard_unpad() is only available at HEAD
diff --git a/docs/notebooks/index.rst b/docs/notebooks/index.rst
index 553597b09..986302dbc 100644
--- a/docs/notebooks/index.rst
+++ b/docs/notebooks/index.rst
@@ -2,7 +2,7 @@ Guided Tour
 =================
 
 To help you get start with the JAX + Flax ecosystem, we suggest that you start with the following notebooks.
-They will walk you through both JAX and Flax basics. 
+They will walk you through both JAX and Flax basics.
 If you are already comfortable with JAX, you can skip that part and start directly with Flax.
 
 .. toctree::
diff --git a/docs/notebooks/linen_intro.ipynb b/docs/notebooks/linen_intro.ipynb
index 53e5ab4b0..821d01ac4 100644
--- a/docs/notebooks/linen_intro.ipynb
+++ b/docs/notebooks/linen_intro.ipynb
@@ -173,7 +173,7 @@
     "id": "ubFTzroGhErh"
    },
    "source": [
-    "We call the `apply` method on the instantiated Module.  If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `apply` with `(variables, *args, rngs=<RNGS>, mutable=<MUTABLEKINDS>, **kwargs)` where \n",
+    "We call the `apply` method on the instantiated Module.  If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `apply` with `(variables, *args, rngs=<RNGS>, mutable=<MUTABLEKINDS>, **kwargs)` where\n",
     " - `<RNGS>` are the optional _call time_ RNGs for things like dropout. For simple Modules this is just a single key, but if your module has multiple __kinds__ of data, it's a dictionary of rng-keys per-kind, e.g. `{'params': key0, 'dropout': key1}` for a Module with dropout layers.\n",
     " - `<MUTABLEKINDS>` is an optional list of names of __kinds__ that are expected to be mutated during the call. e.g. `['batch_stats']` for a layer updating batchnorm statistics.\n",
     "\n",
@@ -392,9 +392,9 @@
     "Flax uses lazy initialization, which allows declared variables to be initialized only at the first site of their use, using whatever shape information is available a the local call site for shape inference.  Once a variable has been initialized, a reference to the data is kept for use in subsequent calls.\n",
     "\n",
     "For declaring parameters that aren't mutated inside the model, but rather by gradient descent, we use the syntax:\n",
-    " \n",
+    "\n",
     " `self.param(parameter_name, parameter_init_fn, *init_args)`\n",
-    " \n",
+    "\n",
     "with arguments:\n",
     " - `parameter_name` just the name, a string\n",
     " - `parameter_init_fn` a function taking an RNG key and a variable number of other arguments, i.e. `fn(rng, *args)`. typically those in `nn.initializers` take an `rng` and a `shape` argument.\n",
@@ -514,7 +514,7 @@
     "  features: int\n",
     "  kernel_init: Callable = nn.initializers.lecun_normal()\n",
     "  bias_init: Callable = nn.initializers.zeros\n",
-    "  \n",
+    "\n",
     "  def setup(self):\n",
     "    self.kernel = self.param('kernel',\n",
     "                             self.kernel_init,\n",
@@ -554,9 +554,9 @@
    },
    "source": [
     "For declaring generally mutable _variables_ that may be mutated inside the model we use the call:\n",
-    " \n",
+    "\n",
     " `self.variable(variable_kind, variable_name, variable_init_fn, *init_args)`\n",
-    " \n",
+    "\n",
     "with arguments:\n",
     " - `variable_kind` the \"kind\" of state this variable is, i.e. the name of the nested-dict collection that this will be stored in inside the top Modules variables.  e.g. `batch_stats` for the moving statistics for a batch norm layer or `cache` for autoregressive cache data.  Note that parameters also have a kind, but they're set to the default `param` kind.\n",
     " - `variable_name` just the name, a string\n",
@@ -750,18 +750,18 @@
     "init_variables = model.init({'params': key2, 'dropout': key3}, x)\n",
     "_, init_params = init_variables.pop('params')\n",
     "\n",
-    "# When calling `apply` with mutable kinds, returns a pair of output, \n",
+    "# When calling `apply` with mutable kinds, returns a pair of output,\n",
     "# mutated_variables.\n",
     "y, mutated_variables = model.apply(\n",
     "    init_variables, x, rngs={'dropout': key4}, mutable=['batch_stats'])\n",
     "\n",
     "# Now we reassemble the full variables from the updates (in a real training\n",
     "# loop, with the updated params from an optimizer).\n",
-    "updated_variables = freeze(dict(params=init_params, \n",
+    "updated_variables = freeze(dict(params=init_params,\n",
     "                                **mutated_variables))\n",
     "\n",
     "print('updated variables:\\n', updated_variables)\n",
-    "print('initialized variable shapes:\\n', \n",
+    "print('initialized variable shapes:\\n',\n",
     "      jax.tree_util.tree_map(jnp.shape, init_variables))\n",
     "print('output:\\n', y)\n",
     "\n",
@@ -900,7 +900,7 @@
    "source": [
     "class RematMLP(nn.Module):\n",
     "  features: Sequence[int]\n",
-    "  # For all transforms, we can annotate a method, or wrap an existing \n",
+    "  # For all transforms, we can annotate a method, or wrap an existing\n",
     "  # Module class. Here we annotate the method.\n",
     "  @nn.remat\n",
     "  @nn.compact\n",
@@ -945,7 +945,7 @@
     "\n",
     "In addition, we provide for each __kind__ of variable it's axis rules:\n",
     "\n",
-    " - `variable_in_axes` - a dict from kinds to a single integer or `None` specifying the input axes to map \n",
+    " - `variable_in_axes` - a dict from kinds to a single integer or `None` specifying the input axes to map\n",
     " - `variable_out_axes` - a dict from kinds to a single integer or `None` specifying the output axes to map\n",
     " - `split_rngs` - a dict from RNG-kinds to a bool, specifying whether to split the rng along the axis.\n",
     "\n",
@@ -994,7 +994,7 @@
     "      attn_weights += bias\n",
     "    norm_dims = tuple(range(attn_weights.ndim // 2, attn_weights.ndim))\n",
     "    attn_weights = jax.nn.softmax(attn_weights, axis=norm_dims)\n",
-    "    attn_weights = nn.Dropout(self.attn_dropout_rate)(attn_weights, \n",
+    "    attn_weights = nn.Dropout(self.attn_dropout_rate)(attn_weights,\n",
     "                                                      deterministic=not self.train)\n",
     "    attn_weights = attn_weights.astype(dtype)\n",
     "\n",
@@ -1098,7 +1098,7 @@
    },
    "source": [
     "Scan allows us to apply `lax.scan` to Modules, including their parameters and mutable variables.  To use it we have to specify how we want each \"kind\" of variable to be transformed.  For scanned variables we specify similar to vmap via in `variable_in_axes`, `variable_out_axes`:\n",
-    " - `nn.broadcast` broadcast the variable kind across the scan steps as a constant \n",
+    " - `nn.broadcast` broadcast the variable kind across the scan steps as a constant\n",
     " - `<axis:int>` scan along `axis` for e.g. unique parameters at each step\n",
     "\n",
     "OR we specify that the variable kind is to be treated like a \"carry\" by passing to the `variable_carry` argument.\n",
@@ -1138,8 +1138,8 @@
     "  @nn.compact\n",
     "  def __call__(self, xs):\n",
     "    dummy_rng = random.PRNGKey(0)\n",
-    "    init_carry = nn.LSTMCell.initialize_carry(dummy_rng, \n",
-    "                                              xs.shape[:1], \n",
+    "    init_carry = nn.LSTMCell.initialize_carry(dummy_rng,\n",
+    "                                              xs.shape[:1],\n",
     "                                              xs.shape[-1])\n",
     "    LSTM = nn.scan(nn.LSTMCell,\n",
     "                   in_axes=1, out_axes=1,\n",
diff --git a/docs/notebooks/linen_intro.md b/docs/notebooks/linen_intro.md
index c86ee2c47..3b1c42bd0 100644
--- a/docs/notebooks/linen_intro.md
+++ b/docs/notebooks/linen_intro.md
@@ -101,7 +101,7 @@ init_variables
 
 +++ {"id": "ubFTzroGhErh"}
 
-We call the `apply` method on the instantiated Module.  If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `apply` with `(variables, *args, rngs=<RNGS>, mutable=<MUTABLEKINDS>, **kwargs)` where 
+We call the `apply` method on the instantiated Module.  If the Module `__call__` method has args `(self, *args, **kwargs)` then we call `apply` with `(variables, *args, rngs=<RNGS>, mutable=<MUTABLEKINDS>, **kwargs)` where
  - `<RNGS>` are the optional _call time_ RNGs for things like dropout. For simple Modules this is just a single key, but if your module has multiple __kinds__ of data, it's a dictionary of rng-keys per-kind, e.g. `{'params': key0, 'dropout': key1}` for a Module with dropout layers.
  - `<MUTABLEKINDS>` is an optional list of names of __kinds__ that are expected to be mutated during the call. e.g. `['batch_stats']` for a layer updating batchnorm statistics.
 
@@ -218,9 +218,9 @@ print('output:\n', y)
 Flax uses lazy initialization, which allows declared variables to be initialized only at the first site of their use, using whatever shape information is available a the local call site for shape inference.  Once a variable has been initialized, a reference to the data is kept for use in subsequent calls.
 
 For declaring parameters that aren't mutated inside the model, but rather by gradient descent, we use the syntax:
- 
+
  `self.param(parameter_name, parameter_init_fn, *init_args)`
- 
+
 with arguments:
  - `parameter_name` just the name, a string
  - `parameter_init_fn` a function taking an RNG key and a variable number of other arguments, i.e. `fn(rng, *args)`. typically those in `nn.initializers` take an `rng` and a `shape` argument.
@@ -280,7 +280,7 @@ class ExplicitDense(nn.Module):
   features: int
   kernel_init: Callable = nn.initializers.lecun_normal()
   bias_init: Callable = nn.initializers.zeros
-  
+
   def setup(self):
     self.kernel = self.param('kernel',
                              self.kernel_init,
@@ -311,9 +311,9 @@ print('output:\n', y)
 +++ {"id": "CJatarOTpByQ"}
 
 For declaring generally mutable _variables_ that may be mutated inside the model we use the call:
- 
+
  `self.variable(variable_kind, variable_name, variable_init_fn, *init_args)`
- 
+
 with arguments:
  - `variable_kind` the "kind" of state this variable is, i.e. the name of the nested-dict collection that this will be stored in inside the top Modules variables.  e.g. `batch_stats` for the moving statistics for a batch norm layer or `cache` for autoregressive cache data.  Note that parameters also have a kind, but they're set to the default `param` kind.
  - `variable_name` just the name, a string
@@ -387,18 +387,18 @@ model = Block(features=3, training=True)
 init_variables = model.init({'params': key2, 'dropout': key3}, x)
 _, init_params = init_variables.pop('params')
 
-# When calling `apply` with mutable kinds, returns a pair of output, 
+# When calling `apply` with mutable kinds, returns a pair of output,
 # mutated_variables.
 y, mutated_variables = model.apply(
     init_variables, x, rngs={'dropout': key4}, mutable=['batch_stats'])
 
 # Now we reassemble the full variables from the updates (in a real training
 # loop, with the updated params from an optimizer).
-updated_variables = freeze(dict(params=init_params, 
+updated_variables = freeze(dict(params=init_params,
                                 **mutated_variables))
 
 print('updated variables:\n', updated_variables)
-print('initialized variable shapes:\n', 
+print('initialized variable shapes:\n',
       jax.tree_util.tree_map(jnp.shape, init_variables))
 print('output:\n', y)
 
@@ -474,7 +474,7 @@ tags: []
 ---
 class RematMLP(nn.Module):
   features: Sequence[int]
-  # For all transforms, we can annotate a method, or wrap an existing 
+  # For all transforms, we can annotate a method, or wrap an existing
   # Module class. Here we annotate the method.
   @nn.remat
   @nn.compact
@@ -510,7 +510,7 @@ You can now `vmap` Modules inside.  The transform has a lot of arguments, they h
 
 In addition, we provide for each __kind__ of variable it's axis rules:
 
- - `variable_in_axes` - a dict from kinds to a single integer or `None` specifying the input axes to map 
+ - `variable_in_axes` - a dict from kinds to a single integer or `None` specifying the input axes to map
  - `variable_out_axes` - a dict from kinds to a single integer or `None` specifying the output axes to map
  - `split_rngs` - a dict from RNG-kinds to a bool, specifying whether to split the rng along the axis.
 
@@ -542,7 +542,7 @@ class RawDotProductAttention(nn.Module):
       attn_weights += bias
     norm_dims = tuple(range(attn_weights.ndim // 2, attn_weights.ndim))
     attn_weights = jax.nn.softmax(attn_weights, axis=norm_dims)
-    attn_weights = nn.Dropout(self.attn_dropout_rate)(attn_weights, 
+    attn_weights = nn.Dropout(self.attn_dropout_rate)(attn_weights,
                                                       deterministic=not self.train)
     attn_weights = attn_weights.astype(dtype)
 
@@ -637,7 +637,7 @@ print('output:\n', y.shape)
 +++ {"id": "8oiRXIC6xQ--"}
 
 Scan allows us to apply `lax.scan` to Modules, including their parameters and mutable variables.  To use it we have to specify how we want each "kind" of variable to be transformed.  For scanned variables we specify similar to vmap via in `variable_in_axes`, `variable_out_axes`:
- - `nn.broadcast` broadcast the variable kind across the scan steps as a constant 
+ - `nn.broadcast` broadcast the variable kind across the scan steps as a constant
  - `<axis:int>` scan along `axis` for e.g. unique parameters at each step
 
 OR we specify that the variable kind is to be treated like a "carry" by passing to the `variable_carry` argument.
@@ -656,8 +656,8 @@ class SimpleScan(nn.Module):
   @nn.compact
   def __call__(self, xs):
     dummy_rng = random.PRNGKey(0)
-    init_carry = nn.LSTMCell.initialize_carry(dummy_rng, 
-                                              xs.shape[:1], 
+    init_carry = nn.LSTMCell.initialize_carry(dummy_rng,
+                                              xs.shape[:1],
                                               xs.shape[-1])
     LSTM = nn.scan(nn.LSTMCell,
                    in_axes=1, out_axes=1,
diff --git a/docs/overview.md b/docs/overview.md
index 46f905ee0..24295bc98 100644
--- a/docs/overview.md
+++ b/docs/overview.md
@@ -24,7 +24,7 @@ JAX that is **designed for flexibility**:
 Try new forms of training by forking an example and by modifying the training
 loop, not by adding features to a framework.
 
-Flax is being developed in close collaboration with the JAX team and 
+Flax is being developed in close collaboration with the JAX team and
 comes with everything you need to start your research, including:
 
 * **Neural network API** (`flax.linen`): Dense, Conv, {Batch|Layer|Group} Norm, Attention, Pooling, {LSTM|GRU} Cell, Dropout
diff --git a/examples/README.md b/examples/README.md
index c028dbac4..0c13b3306 100644
--- a/examples/README.md
+++ b/examples/README.md
@@ -3,10 +3,10 @@
 Each example is designed to be **self-contained and easily forkable**, while
 reproducing relevant results in different areas of machine learning.
 
-As discussed in [#231](https://github.com/google/flax/issues/231), we decided 
-to go for a standard pattern for all examples including the simplest ones 
-(like MNIST). This makes every example a bit more verbose, but once you know 
-one example, you know the structure of all of them. Having unit tests and 
+As discussed in [#231](https://github.com/google/flax/issues/231), we decided
+to go for a standard pattern for all examples including the simplest ones
+(like MNIST). This makes every example a bit more verbose, but once you know
+one example, you know the structure of all of them. Having unit tests and
 integration tests is also very useful when you fork these examples.
 
 For more examples including contributions from the community and other projects currently using Flax see the **[Examples](https://flax.readthedocs.io/en/latest/examples/index.html)** section in the documentation.
diff --git a/examples/imagenet/imagenet_fake_data_benchmark.py b/examples/imagenet/imagenet_fake_data_benchmark.py
index e0dfdee8c..c6116b38f 100644
--- a/examples/imagenet/imagenet_fake_data_benchmark.py
+++ b/examples/imagenet/imagenet_fake_data_benchmark.py
@@ -44,7 +44,7 @@ def test_fake_data(self):
     # Warm-up first so that we are not measuring just compilation.
     with tfds.testing.mock_data(num_examples=1024, data_dir=data_dir):
       train.train_and_evaluate(config, workdir)
-    
+
     start_time = time.time()
     with tfds.testing.mock_data(num_examples=1024, data_dir=data_dir):
       train.train_and_evaluate(config, workdir)
diff --git a/examples/nlp_seq/README.md b/examples/nlp_seq/README.md
index 55d07ca55..57eee7dd7 100644
--- a/examples/nlp_seq/README.md
+++ b/examples/nlp_seq/README.md
@@ -9,7 +9,7 @@ From|ADP the|DT AP|PROPN comes|VBZ this|DT story|NN :|:
 ### Requirements
 * Universal Dependency data sets:  https://universaldependencies.org/#download.
 
-    Download via command line: 
+    Download via command line:
 
     ```
     curl -# -o ud-treebanks-v2.0.tgz https://lindat.mff.cuni.cz/repository/xmlui/bitstream/handle/11234/1-1976/ud-treebanks-v2.0.tgz
@@ -23,7 +23,7 @@ The model should run with other configurations and hardware, but explicitly test
 |:---:|:---:|:---:|:---:|:---:|:---:|
 | Nvidia Titan V (12GB) | 64  |  0.05 | 5:58h | 68.6% | [2022-05-01](https://tensorboard.dev/experiment/F5ULHlyzQlieVJn5PG8mRQ/) |
 
-### Running 
+### Running
 ```
 python train.py --batch_size=64 --model_dir=./ancient_greek \
     --dev=ud-treebanks-v2.0/UD_Ancient_Greek/grc-ud-dev.conllu \
diff --git a/examples/ppo/ppo_lib.py b/examples/ppo/ppo_lib.py
index b17945eb3..32a59264c 100644
--- a/examples/ppo/ppo_lib.py
+++ b/examples/ppo/ppo_lib.py
@@ -288,7 +288,7 @@ def train(
   Returns:
     optimizer: the trained optimizer
   """
-  
+
   game = config.game + 'NoFrameskip-v4'
   simulators = [agent.RemoteSimulator(game)
                 for _ in range(config.num_agents)]
diff --git a/examples/sst2/input_pipeline_test.py b/examples/sst2/input_pipeline_test.py
index 20f1038ad..d591e814e 100644
--- a/examples/sst2/input_pipeline_test.py
+++ b/examples/sst2/input_pipeline_test.py
@@ -81,7 +81,7 @@ def test_batched_dataset_fixed_length(self):
     batch_size = 2
     fixed_pad_length = 77
     for batch in self.dataset.get_batches(
-            batch_size=batch_size, shuffle=False, 
+            batch_size=batch_size, shuffle=False,
             fixed_pad_length=fixed_pad_length).take(1):
       length = batch['token_ids'].numpy().shape[-1]
       self.assertEqual(fixed_pad_length, length)
diff --git a/flax/core/lift.py b/flax/core/lift.py
index ff19daee5..840244765 100644
--- a/flax/core/lift.py
+++ b/flax/core/lift.py
@@ -1165,9 +1165,9 @@ def checkpoint(fn: Callable[..., Any],
       ``pmap``, CSE can defeat the purpose of this decorator. But in some
       settings, like when used inside a ``scan``, this CSE prevention mechanism
       is unnecessary, in which case ``prevent_cse`` can be set to False.
-    static_argnums: Optional, int or sequence of ints, indicates which argument 
-      values on which to specialize for tracing and caching purposes. Specifying 
-      arguments as static can avoid ConcretizationTypeErrors when tracing, but 
+    static_argnums: Optional, int or sequence of ints, indicates which argument
+      values on which to specialize for tracing and caching purposes. Specifying
+      arguments as static can avoid ConcretizationTypeErrors when tracing, but
       at the cost of more retracing overheads.
     policy: Experimental checkpoint policy, see ``jax.checkpoint``.
   Returns:
@@ -1178,7 +1178,7 @@ def inner(scope_fn, repack_fn, variable_groups, rng_groups, *args, **kwargs):
     # add 2 to each static_argnums because we add two initial arguments to rematted
     static_argnums_ = jax.tree_util.tree_map(lambda x: x + 2, static_argnums)
     @functools.partial(jax.remat,
-                       concrete=concrete, static_argnums=static_argnums_, 
+                       concrete=concrete, static_argnums=static_argnums_,
                        prevent_cse=prevent_cse, policy=policy)
     @functools.wraps(fn)
     def rematted(variable_groups, rng_groups, *args, **kwargs):
diff --git a/flax/errors.py b/flax/errors.py
index 0bd7cd596..ed5fb7f67 100644
--- a/flax/errors.py
+++ b/flax/errors.py
@@ -598,14 +598,14 @@ def __init__(self, step, path):
         f' "{path}": No "commit_success.txt" found on this "_gda" directory. '
         'Was its save halted before completion?'
     )
-    
-    
+
+
 class MPARestoreTypeNotMatchError(FlaxError):
   """Make sure the multiprocess array type you use matches your configuration in jax.config.jax_array.
-  
-  If you turned `jax.config.jax_array` on, you should use 
-  `jax.experimental.array.Array` everywhere, instead of using 
-  `GlobalDeviceArray`. Otherwise, avoid using jax.experimental.array 
+
+  If you turned `jax.config.jax_array` on, you should use
+  `jax.experimental.array.Array` everywhere, instead of using
+  `GlobalDeviceArray`. Otherwise, avoid using jax.experimental.array
   to restore your checkpoint.
   """
 
@@ -649,7 +649,7 @@ class BatchDense(nn.Module):
       @nn.compact
       def __call__(self, x):
         return nn.vmap(
-            lambda mdl, x: mdl(x), 
+            lambda mdl, x: mdl(x),
             variable_axes={'params': 0}, split_rngs={'params': True})(nn.Dense(3), x)
 
   """
diff --git a/flax/linen/initializers.py b/flax/linen/initializers.py
index 732f6f15d..791e8cd17 100644
--- a/flax/linen/initializers.py
+++ b/flax/linen/initializers.py
@@ -17,7 +17,7 @@
 # pylint: disable=unused-import
 # re-export initializer functions from jax.nn
 from jax.nn.initializers import constant as constant
-from jax.nn.initializers import delta_orthogonal as delta_orthogonal 
+from jax.nn.initializers import delta_orthogonal as delta_orthogonal
 from jax.nn.initializers import glorot_normal as glorot_normal
 from jax.nn.initializers import glorot_uniform as glorot_uniform
 from jax.nn.initializers import he_normal as he_normal
diff --git a/flax/linen/linear.py b/flax/linen/linear.py
index ca86465ca..47d8abd43 100644
--- a/flax/linen/linear.py
+++ b/flax/linen/linear.py
@@ -468,7 +468,7 @@ def maybe_broadcast(x: Optional[Union[int, Sequence[int]]]) -> (
 
 class Conv(_Conv):
   """Convolution Module wrapping `lax.conv_general_dilated`.
-  
+
   Attributes:
     features: number of convolution filters.
     kernel_size: shape of the convolutional kernel. For 1D convolution,
@@ -511,7 +511,7 @@ def shared_weights(self) -> bool:
 
 class ConvLocal(_Conv):
   """Local convolution Module wrapping `lax.conv_general_dilated_local`.
-  
+
   Attributes:
     features: number of convolution filters.
     kernel_size: shape of the convolutional kernel. For 1D convolution,
diff --git a/flax/linen/module.py b/flax/linen/module.py
index 48c79db8f..c18d6142a 100644
--- a/flax/linen/module.py
+++ b/flax/linen/module.py
@@ -123,7 +123,7 @@ def _module_repr(module: 'Module', num_spaces: int = 4):
 _find_non_lifted_module = re.compile(r'.*\((.*)\)')
 
 def _fix_path_part(part: str):
-  """Fixes a path part by removing transformation name and parenthesis sometimes 
+  """Fixes a path part by removing transformation name and parenthesis sometimes
   inserted by lifted transformations"""
   match = _find_non_lifted_module.match(part)
   if match:
@@ -1537,9 +1537,9 @@ def tabulate(
     **kwargs) -> str:
     """Creates a summary of the Module represented as a table.
 
-    This method has the same signature and internally calls `Module.init`, 
-    but instead of returning the variables, it returns the string summarizing 
-    the Module in a table. `tabulate` uses `jax.eval_shape` to run the forward 
+    This method has the same signature and internally calls `Module.init`,
+    but instead of returning the variables, it returns the string summarizing
+    the Module in a table. `tabulate` uses `jax.eval_shape` to run the forward
     computation without consuming any FLOPs or allocating memory.
 
     Example::
@@ -1602,8 +1602,8 @@ def __call__(self, x):
         name of a single mutable collection. ``list``: A list of names of mutable
         collections. By default all collections except 'intermediates' are
         mutable.
-      console_kwargs: An optional dictionary with additional keyword arguments that 
-        are passed to `rich.console.Console` when rendering the table. Default arguments 
+      console_kwargs: An optional dictionary with additional keyword arguments that
+        are passed to `rich.console.Console` when rendering the table. Default arguments
         are `{'force_terminal': True, 'force_jupyter': False}`.
       **kwargs: keyword arguments to pass to the forward computation.
 
diff --git a/flax/linen/summary.py b/flax/linen/summary.py
index a482109c7..0751489bd 100644
--- a/flax/linen/summary.py
+++ b/flax/linen/summary.py
@@ -130,8 +130,8 @@ def tabulate(
 ) -> Callable[..., str]:
   """Returns a function that creates a summary of the Module represented as a table.
 
-  This function accepts most of the same arguments and internally calls `Module.init`, 
-  except that it returns a function of the form `(*args, **kwargs) -> str` where `*args` 
+  This function accepts most of the same arguments and internally calls `Module.init`,
+  except that it returns a function of the form `(*args, **kwargs) -> str` where `*args`
   and `**kwargs` are passed to `method` (e.g. `__call__`) during the forward pass.
 
   `tabulate` uses `jax.eval_shape` under the hood to run the forward computation without
@@ -157,7 +157,7 @@ def __call__(self, x):
 
   This gives the following output::
 
-                                    Foo Summary                                
+                                    Foo Summary
     ┏━━━━━━━━━┳━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━┳━━━━━━━━━━━━━━━━━━━━━━┓
     ┃ path    ┃ module ┃ inputs        ┃ outputs       ┃ params               ┃
     ┡━━━━━━━━━╇━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━╇━━━━━━━━━━━━━━━━━━━━━━┩
@@ -175,8 +175,8 @@ def __call__(self, x):
     ├─────────┼────────┼───────────────┼───────────────┼──────────────────────┤
     │         │        │               │         Total │ 50 (200 B)           │
     └─────────┴────────┴───────────────┴───────────────┴──────────────────────┘
-                                                                              
-                          Total Parameters: 50 (200 B)                        
+
+                          Total Parameters: 50 (200 B)
 
 
   **Note**: rows order in the table does not represent execution order,
@@ -199,8 +199,8 @@ def __call__(self, x):
     show_repeated: If `True`, repeated calls to the same module will be shown
       in the table, otherwise only the first call will be shown. Default is
       `False`.
-    console_kwargs: An optional dictionary with additional keyword arguments that 
-      are passed to `rich.console.Console` when rendering the table. Default arguments 
+    console_kwargs: An optional dictionary with additional keyword arguments that
+      are passed to `rich.console.Console` when rendering the table. Default arguments
       are `{'force_terminal': True, 'force_jupyter': False}`.
     **kwargs: Additional arguments passed to `Module.init`.
 
@@ -226,7 +226,7 @@ def _get_module_table(
   but returns the Table representation of the Module."""
 
   def _get_table_fn(*args, **kwargs):
-    
+
     with module_lib._tabulate_context():
 
       def _get_variables():
@@ -261,11 +261,11 @@ def _get_variables():
       else:
         module_vars, _ = _get_module_variables(c.path, variables, all_paths)
         counted_vars = module_vars
-        
+
       visited_paths.add(c.path)
       rows.append(
         Row(c.path, c.module_type, c.method, inputs, c.outputs, module_vars, counted_vars))
-        
+
     return Table(module, tuple(collections), rows)
 
   return _get_table_fn
@@ -273,7 +273,7 @@ def _get_variables():
 def _get_module_variables(
   path: Tuple[str, ...], variables: FrozenVariableDict, all_paths: Set[Tuple[str, ...]]
 ) -> Tuple[MutableVariableDict, Any]:
-  """A function that takes a path and variables structure and returns a 
+  """A function that takes a path and variables structure and returns a
   (module_variables, submodule_variables) tuple for that path. _get_module_variables
   uses the `all_paths` set to determine if a variable belongs to a submodule or not."""
   module_variables = _get_path_variables(path, variables)
@@ -291,7 +291,7 @@ def _get_module_variables(
   return module_variables, submodule_variables
 
 def _get_path_variables(path: Tuple[str, ...], variables: FrozenVariableDict) -> MutableVariableDict:
-  """A function that takes a path and a variables structure and returns the variable structure at 
+  """A function that takes a path and a variables structure and returns the variable structure at
   that path."""
   path_variables = {}
 
@@ -305,7 +305,7 @@ def _get_path_variables(path: Tuple[str, ...], variables: FrozenVariableDict) ->
 
     if collection_variables is not None:
       path_variables[collection] = collection_variables.unfreeze()
-    
+
   return path_variables
 
 def _process_inputs(args, kwargs) -> Any:
@@ -327,7 +327,7 @@ def _render_table(table: Table, console_extras: Optional[Mapping[str, Any]]) ->
   console_kwargs = {'force_terminal': True, 'force_jupyter': False}
   if console_extras is not None:
     console_kwargs.update(console_extras)
-    
+
   non_params_cols = 4
   rich_table = rich.table.Table(
       show_header=True,
diff --git a/flax/linen/transforms.py b/flax/linen/transforms.py
index 805b8db41..95f6430b5 100644
--- a/flax/linen/transforms.py
+++ b/flax/linen/transforms.py
@@ -600,9 +600,9 @@ def checkpoint(target: Target,
       ``pmap``, CSE can defeat the purpose of this decorator. But in some
       settings, like when used inside a ``scan``, this CSE prevention mechanism
       is unnecessary, in which case ``prevent_cse`` should be set to False.
-    static_argnums: Optional, int or sequence of ints, indicates which argument 
-      values on which to specialize for tracing and caching purposes. Specifying 
-      arguments as static can avoid ConcretizationTypeErrors when tracing, but 
+    static_argnums: Optional, int or sequence of ints, indicates which argument
+      values on which to specialize for tracing and caching purposes. Specifying
+      arguments as static can avoid ConcretizationTypeErrors when tracing, but
       at the cost of more retracing overheads.
     policy: Experimental checkpoint policy, see ``jax.checkpoint``.
     methods: If `target` is a `Module`, the methods of `Module` to checkpoint.
diff --git a/flax/training/checkpoints.py b/flax/training/checkpoints.py
index 0b90f0b50..5822068bf 100644
--- a/flax/training/checkpoints.py
+++ b/flax/training/checkpoints.py
@@ -490,10 +490,10 @@ def save_checkpoint(ckpt_dir: Union[str, os.PathLike],
                     keep_every_n_steps: Optional[int] = None,
                     async_manager: Optional[AsyncManager] = None) -> str:
   """Save a checkpoint of the model. Suitable for single-host.
-  
+
   In this method, every JAX process saves the checkpoint on its own. Do not
   use it if you have multiple processes and you intend for them to save data
-  to a common directory (e.g., a GCloud bucket). To save multi-process 
+  to a common directory (e.g., a GCloud bucket). To save multi-process
   checkpoints to a shared storage or to save `GlobalDeviceArray`s, use
   `save_checkpoint_multiprocess()` instead.
 
@@ -551,7 +551,7 @@ def save_checkpoint_multiprocess(ckpt_dir: Union[str, os.PathLike],
                                  async_manager: Optional[AsyncManager] = None,
                                  gda_manager: Optional[Any] = None) -> str:
   """Save a checkpoint of the model in multi-process environment.
-  
+
   Use this method to save `GlobalDeviceArray`s, or to save data to a
   common directory. Only process 0 will save the main checkpoint file and
   remove old checkpoint files.
@@ -681,10 +681,10 @@ def restore_checkpoint(
       ckpt_dir must be a directory.
     prefix: str: name prefix of checkpoint files.
     parallel: bool: whether to load seekable checkpoints in parallel, for speed.
-    gda_manager: required if checkpoint contains a multiprocess array 
+    gda_manager: required if checkpoint contains a multiprocess array
       (GlobalDeviceArray or jax Array from pjit). Type should be
       GlobalAsyncCheckpointManager (needs Tensorstore to be imported
-      correctly). Will read the arrays from the separate subdirectory with 
+      correctly). Will read the arrays from the separate subdirectory with
       postfix "_gda".
 
   Returns:
diff --git a/flax/training/early_stopping.py b/flax/training/early_stopping.py
index e3e0bf5ef..d3e81d5e2 100644
--- a/flax/training/early_stopping.py
+++ b/flax/training/early_stopping.py
@@ -34,14 +34,14 @@ class EarlyStopping(struct.PyTreeNode):
       if early_stop.should_stop:
         print('Met early stopping criteria, breaking...')
         break
-  
+
   Attributes:
     min_delta: Minimum delta between updates to be considered an
         improvement.
     patience: Number of steps of no improvement before stopping.
     best_metric: Current best metric value.
     patience_count: Number of steps since last improving update.
-    should_stop: Whether the training loop should stop to avoid 
+    should_stop: Whether the training loop should stop to avoid
         overfitting.
   """
   min_delta: float = 0
@@ -57,7 +57,7 @@ def reset(self):
 
   def update(self, metric):
     """Update the state based on metric.
-    
+
     Returns:
       A pair (has_improved, early_stop), where `has_improved` is True when there
       was an improvement greater than `min_delta` from the previous
diff --git a/flax/training/lr_schedule.py b/flax/training/lr_schedule.py
index d65e6f420..a4c909ed6 100644
--- a/flax/training/lr_schedule.py
+++ b/flax/training/lr_schedule.py
@@ -16,7 +16,7 @@
 
 Note that with `FLIP #1009`_ learning rate schedules in ``flax.training`` are
 **effectively deprecated** in favor of Optax_ schedules. Please refer to
-`Optimizer Schedules`_ for more information. 
+`Optimizer Schedules`_ for more information.
 
 .. _FLIP #1009: https://github.com/google/flax/blob/main/docs/flip/1009-optimizer-api.md
 .. _Optax: https://github.com/deepmind/optax
diff --git a/flax/traverse_util.py b/flax/traverse_util.py
index 2531f4028..6b9d03471 100644
--- a/flax/traverse_util.py
+++ b/flax/traverse_util.py
@@ -168,7 +168,7 @@ def path_aware_map(
   the path to each leaf into account.
 
   Example::
-    
+
     >>> import jax.numpy as jnp
     >>> from flax import traverse_util
     ...
@@ -181,7 +181,7 @@ def path_aware_map(
     f: A callable that takes in ``(path, value)`` arguments and maps them
       to a new value. Here ``path`` is a tuple of strings.
     nested_dict: A nested dictionary structure.
-  
+
   Returns:
     A new nested dictionary structure with the mapped values.
   """
diff --git a/pytest.ini b/pytest.ini
index f18222255..cfbfc76d0 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -16,7 +16,7 @@ filterwarnings =
     ignore:the imp module is deprecated.*:DeprecationWarning
 # traverse_util.Traversal will be removed soon.
     ignore:`flax.traverse_util.Traversal` will be deprecated.*:DeprecationWarning
-# TODO: Will revisit all the deprecation warnings next week. 
+# TODO: Will revisit all the deprecation warnings next week.
     ignore:jax.tree_.*:FutureWarning
 # TF internal deprecation warnings.
     ignore:Call to deprecated create function .*Descriptor.*:DeprecationWarning
diff --git a/setup.cfg b/setup.cfg
index 1faa578d1..d18102719 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,4 +1,4 @@
 [report]
-exclude_lines = 
+exclude_lines =
 	@abc.abstractmethod
 	raise NotImplementedError
diff --git a/tests/checkpoints_test.py b/tests/checkpoints_test.py
index de3b2ba59..e531dd32f 100644
--- a/tests/checkpoints_test.py
+++ b/tests/checkpoints_test.py
@@ -307,22 +307,22 @@ def test_last_checkpoint(self):
     with gfile.GFile(os.path.join(tmp_dir, 'test_tmp'), 'w') as f:
         f.write('test_tmp')
     gfile.makedirs(os.path.join(tmp_dir, 'test_tmp_gda'))
-    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'), 
+    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'),
                      None)
-                     
+
     with gfile.GFile(os.path.join(tmp_dir, 'test_0'), 'w') as f:
         f.write('test_0')
     gfile.makedirs(os.path.join(tmp_dir, 'test_0_gda'))
-    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'), 
+    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'),
                      os.path.join(tmp_dir, 'test_0'))
-    
+
     with gfile.GFile(os.path.join(tmp_dir, 'test_10'), 'w') as f:
         f.write('test_10')
-    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'), 
+    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'test_'),
                      os.path.join(tmp_dir, 'test_10'))
-    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'ckpt_'), 
+    self.assertEqual(checkpoints.latest_checkpoint(tmp_dir, 'ckpt_'),
                      None)
-    
+
   @parameterized.parameters({'jax_array_config': True},
                             {'jax_array_config': False})
   def test_jax_array(self, jax_array_config):
@@ -342,7 +342,7 @@ def test_jax_array(self, jax_array_config):
         tmp_dir, test_object0, prefix='test_')
     jtu.check_eq(new_object, {'a': np.ones(3), 'b': np.arange(3, 6)})
 
-  
+
   def test_convert_pre_linen(self):
     params = checkpoints.convert_pre_linen({
         'mod_0': {
diff --git a/tests/core/core_scope_test.py b/tests/core/core_scope_test.py
index e52b93900..177572716 100644
--- a/tests/core/core_scope_test.py
+++ b/tests/core/core_scope_test.py
@@ -167,7 +167,7 @@ def test_rng_check_w_old_and_new_keys(self):
       self.assertTrue(scope._is_valid_rng(random.PRNGKey(0)))
     finally:
       jax_config.update('jax_enable_custom_prng', old_setting)
-  
+
   def test_jax_leak_detector(self):
     with jax.check_tracer_leaks(True):
       def f(scope):
@@ -199,7 +199,7 @@ def test_empty_col_error(self):
     root = Scope({'state': {}})
     with self.assertRaises(errors.ScopeCollectionNotFound):
       root.variable('state', 'test', jnp.zeros, ())
-  
+
   def test_variable_no_init(self):
     root = Scope({}, mutable='state')
     with self.assertRaises(errors.ScopeCollectionNotFound):
@@ -209,7 +209,7 @@ def test_variable_no_init(self):
     self.assertEqual(abc.value, 1)
     with self.assertRaises(errors.ScopeVariableNotFoundError):
       root.variable('state', 'test')
-  
+
   def test_variable_alias(self):
     scope = Scope({}, mutable='state')
     subscope = scope.push(name="a")
diff --git a/tests/core/design/core_attention_test.py b/tests/core/design/core_attention_test.py
index 39c692597..99a09b959 100644
--- a/tests/core/design/core_attention_test.py
+++ b/tests/core/design/core_attention_test.py
@@ -140,7 +140,7 @@ class AttentionTest(absltest.TestCase):
   def test_attention(self):
     inputs = jnp.ones((2, 7, 16))
     model = partial(
-        multi_head_dot_product_attention, 
+        multi_head_dot_product_attention,
         num_heads=2, batch_axes=(0,),
         attn_fn=with_dropout(softmax_attn, 0.1, deterministic=False))
 
diff --git a/tests/core/design/core_auto_encoder_test.py b/tests/core/design/core_auto_encoder_test.py
index 6a4c8341f..c97021fe2 100644
--- a/tests/core/design/core_auto_encoder_test.py
+++ b/tests/core/design/core_auto_encoder_test.py
@@ -57,7 +57,7 @@ def module_method(fn, name=None):
     name = fn.__name__ if hasattr(fn, '__name__') else None
 
   def wrapper(self, *args, **kwargs):
-    scope = self.scope.rewound()  
+    scope = self.scope.rewound()
     mod_fn = lambda scope: fn(self, scope, *args, **kwargs)
     return scope.child(mod_fn, name)()
   return wrapper
@@ -135,7 +135,7 @@ def test_auto_encoder_with_scope(self):
             'out': {'kernel': (3, 4), 'bias': (4,)},
         },
     })
-  
+
   def test_auto_encoder_bind_method(self):
     ae = lambda scope, x: AutoEncoder3.create(scope, latents=2, features=4, hidden=3)(x)
     x = jnp.ones((1, 4))
diff --git a/tests/core/design/core_dense_test.py b/tests/core/design/core_dense_test.py
index 5e8c1ca2c..b3cc390c8 100644
--- a/tests/core/design/core_dense_test.py
+++ b/tests/core/design/core_dense_test.py
@@ -129,7 +129,7 @@ def test_explicit_dense(self):
         'dense_0': ExplicitDense((4, 3), (3,)),
         'dense_1': ExplicitDense((3, 1), (1,))
     })
-  
+
   def test_semi_explicit_dense(self):
     x = jnp.ones((1, 4))
     y, variables = init(semi_explicit_mlp)(random.PRNGKey(0), x)
diff --git a/tests/core/design/core_scan_test.py b/tests/core/design/core_scan_test.py
index 85575a923..3e6031a78 100644
--- a/tests/core/design/core_scan_test.py
+++ b/tests/core/design/core_scan_test.py
@@ -61,11 +61,11 @@ def test_scan_unshared_params(self):
     self.assertEqual(param_shapes, {
       'dense_0': {'kernel': (2, 4, 1), 'bias': (2, 1)},
     })
-    
+
     self.assertNotEqual(y[0], y[1])
     k1, k2 = variables['params']['dense_0']['kernel']
     self.assertFalse(jnp.allclose(k1, k2))
-  
+
   def test_scan_shared_params(self):
     x = random.normal(random.PRNGKey(0), (1, 4))
     x = jnp.concatenate([x, x], 0)
@@ -77,7 +77,7 @@ def test_scan_shared_params(self):
     self.assertEqual(param_shapes, {
       'dense_0': {'kernel': (4, 1), 'bias': (1,)},
     })
-    
+
     self.assertEqual(y[0], y[1])
 
 
diff --git a/tests/early_stopping_test.py b/tests/early_stopping_test.py
index bfa159bb8..685cb4627 100644
--- a/tests/early_stopping_test.py
+++ b/tests/early_stopping_test.py
@@ -29,9 +29,9 @@
 class EarlyStoppingTests(absltest.TestCase):
 
   def test_update(self):
-    es = early_stopping.EarlyStopping(min_delta=0, 
+    es = early_stopping.EarlyStopping(min_delta=0,
                                       patience=0)
-    
+
     for i in range(2):
       improve_steps = 0
       for step in range(10):
@@ -48,9 +48,9 @@ def test_update(self):
       es = es.reset()  # ensure object is reusable if reset.
 
   def test_patience(self):
-    es = early_stopping.EarlyStopping(min_delta=0, 
+    es = early_stopping.EarlyStopping(min_delta=0,
                                       patience=0)
-    patient_es = early_stopping.EarlyStopping(min_delta=0, 
+    patient_es = early_stopping.EarlyStopping(min_delta=0,
                                               patience=6)
     for step in range(10):
       metric = 1.
@@ -69,11 +69,11 @@ def test_patience(self):
     self.assertEqual(patient_step, 7)
 
   def test_delta(self):
-    es = early_stopping.EarlyStopping(min_delta=0, 
+    es = early_stopping.EarlyStopping(min_delta=0,
                                       patience=0)
-    delta_es = early_stopping.EarlyStopping(min_delta=1e-3, 
+    delta_es = early_stopping.EarlyStopping(min_delta=1e-3,
                                             patience=0)
-    delta_patient_es = early_stopping.EarlyStopping(min_delta=1e-3, 
+    delta_patient_es = early_stopping.EarlyStopping(min_delta=1e-3,
                                                     patience=1)
     metric = 1.
     for step in range(100):
@@ -90,10 +90,10 @@ def test_delta(self):
       did_improve, delta_es = delta_es.update(metric)
       if delta_es.should_stop:
         break
-      
+
     self.assertEqual(step, 1)
 
-    metrics = [0.01, 0.005, 0.0033, 0.0025, 0.002, 
+    metrics = [0.01, 0.005, 0.0033, 0.0025, 0.002,
                0.0017, 0.0014, 0.0012, 0.0011, 0.001]
     improvement_steps = 0
     for step in range(10):
@@ -103,7 +103,7 @@ def test_delta(self):
         improvement_steps += 1
       if delta_patient_es.should_stop:
         break
-      
+
     self.assertEqual(improvement_steps, 4)  # steps 0, 1, 2, 4
     self.assertEqual(step, 6)
 
diff --git a/tests/linen/linen_module_test.py b/tests/linen/linen_module_test.py
index e27f55d08..6ef79157c 100644
--- a/tests/linen/linen_module_test.py
+++ b/tests/linen/linen_module_test.py
@@ -1428,7 +1428,7 @@ def __call__(self, x):
         x = 4 * x
         x = self.perturb('after_multiply', x)
         return x
-    
+
     def loss(params, perturbations, inputs, targets):
       variables = {'params': params, 'perturbations': perturbations}
       preds = Foo().apply(variables, inputs)
diff --git a/tests/linen/linen_transforms_test.py b/tests/linen/linen_transforms_test.py
index 292a2bb77..6e9a60011 100644
--- a/tests/linen/linen_transforms_test.py
+++ b/tests/linen/linen_transforms_test.py
@@ -147,7 +147,7 @@ def __call__(self, input, apply_relu : bool = False):
 
   def test_remat_static_argnums(self):
     test = self
-    
+
     class Foo(nn.Module):
       train_is_static: bool
 
@@ -157,7 +157,7 @@ def __call__(self, inputs, train: bool):
           test.assertTrue(isinstance(train, bool))
         else:
           test.assertTrue(isinstance(train, jnp.ndarray))
-        
+
         return nn.Dense(3, use_bias=False)(inputs)
 
     # set train as a static argument
@@ -176,7 +176,7 @@ def __call__(self, inputs, train: bool):
     variables = foo.init(random.PRNGKey(0), x, True)
     y = foo.apply(variables, x, False)
     self.assertEqual(y.shape, (1, 3))
-  
+
   def test_remat_decorator_static_argnums(self):
     test = self
 
@@ -185,7 +185,7 @@ class FooTrainStatic(nn.Module):
       @nn.compact
       def __call__(self, inputs, train: bool):
         test.assertTrue(isinstance(train, bool))
-        
+
         return nn.Dense(3, use_bias=False)(inputs)
 
     # set train as a static argument
@@ -201,7 +201,7 @@ class FooTrainDynamic(nn.Module):
       @nn.compact
       def __call__(self, inputs, train: bool):
         test.assertTrue(isinstance(train, jnp.ndarray))
-        
+
         return nn.Dense(3, use_bias=False)(inputs)
 
     # set train as a non-static arguments
@@ -210,7 +210,7 @@ def __call__(self, inputs, train: bool):
     variables = foo.init(random.PRNGKey(0), x, True)
     y = foo.apply(variables, x, False)
     self.assertEqual(y.shape, (1, 3))
-    
+
 
   def test_vmap(self):
     key1, key2 = random.split(random.PRNGKey(3), 2)
diff --git a/tests/linen/summary_test.py b/tests/linen/summary_test.py
index 5b80c819b..a0335b47b 100644
--- a/tests/linen/summary_test.py
+++ b/tests/linen/summary_test.py
@@ -52,7 +52,7 @@ def block_method(self, x: Array, training: bool) -> Array:
     x = self.dropout(x, deterministic=not training)
     x = nn.relu(x)
     return x
-  
+
   def __call__(self, x: Array, training: bool) -> Array:
     x = self.conv(x)
 
@@ -71,7 +71,7 @@ def setup(self) -> None:
     self.block1 = ConvBlock(32, [3, 3], test_sow=self.test_sow)
     self.block2 = ConvBlock(64, [3, 3], test_sow=self.test_sow)
     self.dense = nn.Dense(10)
-  
+
   def cnn_method(self, x: Array, training: bool) -> Array:
     x = self.block1.block_method(x, training=training)
     x = self.block2.block_method(x, training=training)
@@ -109,9 +109,9 @@ def test_module_summary(self):
     x = jnp.ones((batch_size, 28, 28, 1))
     module = CNN(test_sow=False)
 
-    table = summary._get_module_table(module, depth=None, show_repeated=True)( 
+    table = summary._get_module_table(module, depth=None, show_repeated=True)(
       {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
-      x, training=True, mutable=True, 
+      x, training=True, mutable=True,
     )
     # get values for inputs and outputs from their _ValueRepresentation
     for row in table:
@@ -173,7 +173,7 @@ def test_module_summary(self):
         row.module_variables,
         row.counted_variables,
       )
-  
+
   def test_module_summary_with_depth(self):
     """
     This test creates a Table using `module_summary` set the `depth` argument to `1`,
@@ -186,7 +186,7 @@ def test_module_summary_with_depth(self):
 
     table = summary._get_module_table(module, depth=1, show_repeated=True)(
       {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
-      x, training=True, mutable=True, 
+      x, training=True, mutable=True,
     )
     # get values for inputs and outputs from their _ValueRepresentation
     for row in table:
@@ -230,10 +230,10 @@ def test_module_summary_with_depth(self):
     self.assertEqual(table[0].module_variables, table[0].counted_variables)
     self.assertEqual(table[3].module_variables, table[3].counted_variables)
 
-  
+
   def test_tabulate(self):
     """
-    This test creates a string representation of a Module using `Module.tabulate` 
+    This test creates a string representation of a Module using `Module.tabulate`
     and checks that it matches the expected output given the CNN model defined in `_get_tabulate_cnn`.
     """
     batch_size = 32
@@ -242,8 +242,8 @@ def test_tabulate(self):
     module = CNN(test_sow=False)
 
     module_repr = module.tabulate(
-        {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)}, 
-        x, 
+        {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
+        x,
         training=True,
         console_kwargs=CONSOLE_TEST_KWARGS,
     )
@@ -263,7 +263,7 @@ def test_tabulate(self):
     self.assertIn("inputs", lines[3])
     self.assertIn("outputs", lines[3])
     self.assertIn("params", lines[3])
-    self.assertIn("batch_stats", lines[3]) 
+    self.assertIn("batch_stats", lines[3])
 
     # collection counts
     self.assertIn("Total", lines[-6])
@@ -286,15 +286,15 @@ def test_tabulate_with_sow(self):
     module = CNN(test_sow=True)
 
     module_repr = module.tabulate(
-      {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)}, 
-      x, 
+      {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
+      x,
       training=True,
       console_kwargs=CONSOLE_TEST_KWARGS,
     )
 
     self.assertIn("intermediates", module_repr)
     self.assertIn("INTERM", module_repr)
-  
+
   def test_tabulate_with_method(self):
 
     batch_size = 32
@@ -303,8 +303,8 @@ def test_tabulate_with_method(self):
     module = CNN(test_sow=False)
 
     module_repr = module.tabulate(
-      {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)}, 
-      x, 
+      {"dropout":random.PRNGKey(0), "params": random.PRNGKey(1)},
+      x,
       training=True,
       method=CNN.cnn_method,
       console_kwargs=CONSOLE_TEST_KWARGS,
@@ -315,7 +315,7 @@ def test_tabulate_with_method(self):
 
   def test_tabulate_function(self):
     """
-    This test creates a string representation of a Module using `Module.tabulate` 
+    This test creates a string representation of a Module using `Module.tabulate`
     and checks that it matches the expected output given the CNN model defined in `_get_tabulate_cnn`.
     """
     batch_size = 32
@@ -384,9 +384,9 @@ def __call__(self, x):
     with jax.check_tracer_leaks(True):
       module_repr = lstm.tabulate(
         random.PRNGKey(0),
-        x=jnp.ones((32, 128, 64)), 
+        x=jnp.ones((32, 128, 64)),
         console_kwargs=CONSOLE_TEST_KWARGS)
-    
+
     lines = module_repr.splitlines()
 
     self.assertIn("LSTM", lines[5])
@@ -394,7 +394,7 @@ def __call__(self, x):
     self.assertIn("LSTMCell", lines[9])
     self.assertIn("ScanLSTM/ii", lines[13])
     self.assertIn("Dense", lines[13])
-  
+
   def test_lifted_transform_no_rename(self):
     class LSTM(nn.Module):
       batch_size: int
@@ -420,7 +420,7 @@ def __call__(self, x):
     with jax.check_tracer_leaks(True):
       module_repr = lstm.tabulate(
         random.PRNGKey(0),
-        x=jnp.ones((32, 128, 64)), 
+        x=jnp.ones((32, 128, 64)),
         console_kwargs=CONSOLE_TEST_KWARGS)
 
     lines = module_repr.splitlines()
@@ -452,7 +452,7 @@ def __call__(self, x):
 
     x = jnp.ones((4, 28, 28, 32))
     module_repr = CNN().tabulate(
-      jax.random.PRNGKey(0), 
+      jax.random.PRNGKey(0),
       x=x,
       show_repeated=True,
       console_kwargs=CONSOLE_TEST_KWARGS)
@@ -473,7 +473,7 @@ def __call__(self, x):
     self.assertNotIn("mean", lines[25])
     self.assertNotIn("bias", lines[25])
     self.assertIn("ConvBlock_0/Dropout_0", lines[27])
-    
+
     # third call
     self.assertIn("ConvBlock_0/Conv_0", lines[31])
     self.assertNotIn("bias", lines[31])
@@ -493,7 +493,7 @@ def __call__(self):
     lines = module_repr.splitlines()
 
     self.assertRegex(lines[5], r'|\s*|\s*EmptyInput\s*|\s*|\s*1\s*|')
-  
+
   def test_numpy_scalar(self):
     class Submodule(nn.Module):
       def __call__(self, x):
diff --git a/tests/tensorboard_test.py b/tests/tensorboard_test.py
index 99c2e3f25..40fea447c 100644
--- a/tests/tensorboard_test.py
+++ b/tests/tensorboard_test.py
@@ -290,7 +290,7 @@ def test_flatten_dict(self):
         "None": None,
         "List": [1, 2, 3],
         "Tuple": (1, 2, 3),
-        "Complex": complex("1+1j"), 
+        "Complex": complex("1+1j"),
         "np.complex_": np.complex_("1+1j"),
         "Bool": True,
         "Int": 1,
@@ -302,7 +302,7 @@ def test_flatten_dict(self):
         "np.character": np.str_("test")
       }
     }
-    
+
     result_hparams = _flatten_dict(input_hparams)
 
     expected_hparams={
diff --git a/tests/traverse_util_test.py b/tests/traverse_util_test.py
index 0e8cf8e50..ebc43168f 100644
--- a/tests/traverse_util_test.py
+++ b/tests/traverse_util_test.py
@@ -260,9 +260,9 @@ def test_path_value(self):
     params_in = {'a': {'b': 10, 'c': 2}}
     params_out = traverse_util.path_aware_map(
       lambda path, x: x + 1 if 'b' in path else -x, params_in)
-    
+
     self.assertEqual(params_out, {'a': {'b': 11, 'c': -2}})
-    
+
   def test_path_aware_map_with_multi_transform(self):
     params = {'linear_1': {'w': jnp.zeros((5, 6)), 'b': jnp.zeros(5)},
             'linear_2': {'w': jnp.zeros((6, 1)), 'b': jnp.zeros(1)}}
@@ -275,13 +275,13 @@ def test_path_aware_map_with_multi_transform(self):
     state = tx.init(params)
     updates, new_state = tx.update(gradients, state, params)
     new_params = optax.apply_updates(params, updates)
-    
+
 
     self.assertTrue(np.allclose(new_params['linear_1']['b'], params['linear_1']['b']))
     self.assertTrue(np.allclose(new_params['linear_2']['b'], params['linear_2']['b']))
     self.assertFalse(np.allclose(new_params['linear_1']['w'], params['linear_1']['w']))
     self.assertFalse(np.allclose(new_params['linear_2']['w'], params['linear_2']['w']))
-  
+
   def test_path_aware_map_with_masked(self):
     params = {'linear_1': {'w': jnp.zeros((5, 6)), 'b': jnp.zeros(5)},
             'linear_2': {'w': jnp.zeros((6, 1)), 'b': jnp.zeros(1)}}
@@ -293,7 +293,7 @@ def test_path_aware_map_with_masked(self):
     state = tx.init(params)
     updates, new_state = tx.update(gradients, state, params)
     new_params = optax.apply_updates(params, updates)
-    
+
 
     self.assertTrue(np.allclose(new_params['linear_1']['b'], gradients['linear_1']['b']))
     self.assertTrue(np.allclose(new_params['linear_2']['b'], gradients['linear_2']['b']))
@@ -304,7 +304,7 @@ def test_path_aware_map_with_empty_nodes(self):
     params_in = {'a': {'b': 10, 'c': 2}, 'b': {}}
     params_out = traverse_util.path_aware_map(
       lambda path, x: x + 1 if 'b' in path else -x, params_in)
-    
+
     self.assertEqual(params_out, {'a': {'b': 11, 'c': -2}, 'b': {}})