Merge pull request #2732 from levskaya:refixci

PiperOrigin-RevId: 495997966

Author: Flax Authors

.github/workflows/build.yml

@@ -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.8', '3.9', '3.10']