diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml
index 932176e78..05743081b 100644
--- a/.github/workflows/build.yml
+++ b/.github/workflows/build.yml
@@ -28,7 +28,7 @@ jobs:
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python 3.8
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v4
         with:
           python-version: 3.7
       - uses: pre-commit/action@v2.0.3
@@ -36,7 +36,7 @@ jobs:
     name: Check commit count
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     # We allow at most 5 commits in a branch to ensure our CI doesn't break.
     - name: Check commit count in PR
       if: always()
@@ -65,9 +65,9 @@ jobs:
       matrix:
         python-version: [3.7, 3.8, 3.9]
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
     - name: Install standalone dependencies only
@@ -79,7 +79,7 @@ jobs:
   tests:
     name: Run Tests
     needs: [cancel-previous, pre-commit, commit-count, test-import]
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-20.04-16core
     strategy:
       matrix:
         python-version: [3.7, 3.8, 3.9]
@@ -90,24 +90,40 @@ jobs:
           - test-type: pytype
             python-version: 3.8
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v1
+      id: setup_python
+      uses: actions/setup-python@v4
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Install dependencies
+    - name: Get week and year
+      id: week_year
+      run: echo "DATE=$(date +%U-%Y)" >> $GITHUB_OUTPUT
+    - name: Cached virtual environment
+      id: venv_cache
+      uses: actions/cache@v3
+      with:
+        path: venv
+        key: pip-${{ steps.setup_python.outputs.python-version }}-${{ steps.week_year.outputs.DATE }}-${{ hashFiles('**/requirements.txt', 'setup.py') }}
+    - name: Install Dependencies for cache
+      if: steps.venv_cache.outputs.cache-hit != 'true'
       run: |
-        pip install .
-        pip install .[testing]
-        pip install -r docs/requirements.txt
+        if [ -d "venv" ]; then rm -rf venv; fi
+        python3 -m venv venv
+        venv/bin/python3 -m pip install .
+        venv/bin/python3 -m pip install .[testing]
+        venv/bin/python3 -m pip install -r docs/requirements.txt
+    - name: Install Flax
+      run: |
+        venv/bin/python3 -m pip install -e .
     - name: Test with ${{ matrix.test-type }}
       run: |
         if [[ "${{ matrix.test-type }}" == "doctest" ]]; then
-          tests/run_all_tests.sh --no-pytest --no-pytype
+          tests/run_all_tests.sh --no-pytest --no-pytype --use-venv
         elif [[ "${{ matrix.test-type }}" == "pytest" ]]; then
-          tests/run_all_tests.sh --no-doctest --no-pytype --with-cov
+          tests/run_all_tests.sh --no-doctest --no-pytype --with-cov --use-venv
         elif [[ "${{ matrix.test-type }}" == "pytype" ]]; then
-          tests/run_all_tests.sh --no-doctest --no-pytest
+          tests/run_all_tests.sh --no-doctest --no-pytest --use-venv
         else
           echo "Unknown test type: ${{ matrix.test-type }}"
           exit 1
diff --git a/tests/run_all_tests.sh b/tests/run_all_tests.sh
index 3af60f876..8164b8289 100755
--- a/tests/run_all_tests.sh
+++ b/tests/run_all_tests.sh
@@ -4,6 +4,7 @@ PYTEST_OPTS=
 RUN_DOCTEST=true
 RUN_PYTEST=true
 RUN_PYTYPE=true
+GH_VENV=false
 
 for flag in "$@"; do
 case $flag in
@@ -24,6 +25,9 @@ case $flag in
   --no-pytype)
   RUN_PYTYPE=false
   ;;
+  --use-venv)
+  GH_VENV=true
+  ;;
   *)
   echo "Unknown flag: $flag"
   exit 1
@@ -31,11 +35,18 @@ case $flag in
 esac
 done
 
+# Activate cached virtual env for github CI
+if $GH_VENV; then
+  source $(dirname "$0")/../venv/bin/activate
+fi
+
 echo "====== test config ======="
 echo "PYTEST_OPTS: $PYTEST_OPTS"
 echo "RUN_DOCTEST: $RUN_DOCTEST"
 echo "RUN_PYTEST: $RUN_PYTEST"
 echo "RUN_PYTYPE: $RUN_PYTYPE"
+echo "GH_VENV: $GH_VENV"
+echo "WHICH PYTHON: $(which python)"
 echo "jax: $(python -c 'import jax; print(jax.__version__)')"
 echo "flax: $(python -c 'import flax; print(flax.__version__)')"
 echo "=========================="
@@ -99,14 +110,14 @@ fi
 if $RUN_PYTYPE; then
   echo "=== RUNNING PYTYPE ==="
   # Validate types in library code.
-  pytype --config pytype.cfg flax/
+  pytype --jobs auto --config pytype.cfg flax/
 
   # Validate types in examples.
   for egd in $(find examples -maxdepth 1 -mindepth 1 -type d); do
       # use cd to make sure pytpe cache lives in example dir and doesn't name clash
       # use *.py to avoid importing configs as a top-level import which leads to import errors
       # because config files use relative imports (e.g. from config import ...).
-      (cd $egd ; pytype --config ../../pytype.cfg "*.py")
+      (cd $egd ; pytype --jobs auto --config ../../pytype.cfg "*.py")
   done
 fi