Merge branch 'main' into fix_dropout

Author: chiamp

.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:

.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

.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 @@
                 }
               }
             }
-          } 
+          }
         }
       }
     }

.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) {

.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.

.github/workflows/build.yml

@@ -13,16 +13,21 @@ on:
       - main
 
 jobs:
-  pre-commit:
-    name: Test pre-commit hooks
+  cancel-previous:
+    name: Cancel Previous Runs
     runs-on: ubuntu-latest
-    timeout-minutes: 5
     steps:
       - name: Cancel previous
-        uses: styfle/[email protected]
+        uses: styfle/[email protected]
+        if: ${{github.ref != 'refs/head/main'}}
         with:
           access_token: ${{ github.token }}
-        if: ${{github.ref != 'refs/head/main'}}
+          workflow_id: ${{ github.event.workflow.id }}
+          all_but_latest: true
+  pre-commit:
+    name: Test pre-commit hooks
+    runs-on: ubuntu-latest
+    steps:
       - uses: actions/checkout@v3
       - name: Set up Python 3.8
         uses: actions/setup-python@v3
@@ -33,17 +38,13 @@ jobs:
     name: Check commit count
     runs-on: ubuntu-latest
     steps:
-    - name: Cancel previous
-      uses: styfle/[email protected]
-      with:
-        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
@@ -59,8 +60,27 @@ jobs:
           echo "See $url for help on how to resolve this."
           exit 1
         fi
+  test-import:
+    name: Test import standalone
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [3.7, 3.8, 3.9]
+    steps:
+    - uses: actions/checkout@v2
+    - name: Set up Python ${{ matrix.python-version }}
+      uses: actions/setup-python@v1
+      with:
+        python-version: ${{ matrix.python-version }}
+    - name: Install standalone dependencies only
+      run: |
+        pip install .
+    - name: Test importing Flax
+      run: |
+        python -c "import flax"
   tests:
     name: Run Tests
+    needs: [cancel-previous, pre-commit, commit-count, test-import]
     runs-on: ubuntu-latest
     strategy:
       matrix:
@@ -72,10 +92,6 @@ jobs:
           - test-type: pytype
             python-version: 3.8
     steps:
-    - name: Cancel previous
-      uses: styfle/[email protected]
-      with:
-        access_token: ${{ github.token }}          
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v1
@@ -121,23 +137,3 @@ jobs:
            "description": "'$status'",
            "context": "github-actions/Build"
            }'
-  test-import:
-    name: Test import standalone
-    runs-on: ubuntu-latest
-    strategy:
-      matrix:
-        python-version: [3.7, 3.8, 3.9]
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
-      with:
-        python-version: ${{ matrix.python-version }}
-    - name: Install standalone dependencies only
-      run: |
-        pip install .
-    - name: Test importing Flax
-      run: |
-        python -c "import flax"
-
-

.pre-commit-config.yaml

@@ -13,3 +13,7 @@ repos:
   hooks:
   - id: jupytext
     args: [--sync]
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: v4.3.0
+  hooks:
+  - id: trailing-whitespace

.readthedocs.yml

@@ -12,6 +12,8 @@ sphinx:
 # Optionally build your docs in additional formats such as PDF and ePub
 formats:
   - htmlzip
+  - epub
+  - pdf
 
 # Optionally set the version of Python and requirements required to build your docs
 python:

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

README.md

@@ -87,10 +87,11 @@ 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).
+[guides](https://flax.readthedocs.io/en/latest/guides/index.html) and
+[advanced topics documentation](https://flax.readthedocs.io/en/latest/advanced_topics/index.html).
 
 ```py
 from typing import Sequence
@@ -174,11 +175,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

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",

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:

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}']

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

docs/_templates/autosummary/flax_module.rst

@@ -7,8 +7,8 @@
 
    {% block methods %}
 
-   .. automethod:: __call__ 
-   
+   .. automethod:: __call__
+
    {% if methods %}
    .. rubric:: Methods