Add pyrefly type checking: pyrefly.toml config, CI integration, test runner support

Agent-Logs-Url: https://github.com/vfdev-5/flax/sessions/d862656c-0c03-46b4-971e-55ff771b19b9

Author: vfdev-5

.github/workflows/flax_test.yml

@@ -98,6 +98,9 @@ jobs:
           - python-version: '3.12'
             test-type: mypy
             jax-version: 'newest'
+          - python-version: '3.12'
+            test-type: pyrefly
+            jax-version: 'newest'
     steps:
     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
     - name: Setup uv
@@ -128,6 +131,8 @@ jobs:
           uv run --no-sync tests/run_all_tests.sh --only-pytype
         elif [[ "${{ matrix.test-type }}" == "mypy" ]]; then
           uv run --no-sync tests/run_all_tests.sh --only-mypy
+        elif [[ "${{ matrix.test-type }}" == "pyrefly" ]]; then
+          uv run --no-sync tests/run_all_tests.sh --only-pyrefly
         else
           echo "Unknown test type: ${{ matrix.test-type }}"
           exit 1

flax/linen/linear.py

@@ -189,6 +189,7 @@ def bias_init_wrap(rng, shape, dtype=jnp.float32):
     inputs, kernel, bias = self.promote_dtype(
       inputs, kernel, bias, dtype=self.dtype
     )
+    assert inputs is not None and kernel is not None
 
     if self.dot_general_cls is not None:
       dot_general = self.dot_general_cls()

pyproject.toml

@@ -45,6 +45,7 @@ testing = [
     "jraph>=0.0.6dev0",
     "ml-collections",
     "mypy",
+    "pyrefly",
     "opencv-python",
     # Set protobuf version to prevent error in
     # examples/mnist/train_test.py::TrainTest::test_train_and_evaluate

pyrefly.toml

@@ -0,0 +1,26 @@
+# Pyrefly configuration - migrated from mypy
+# Only type-check flax/linen/linear.py for now; expand as issues are resolved.
+project-includes = ["flax/linen/linear.py"]
+
+preset = "legacy"
+ignore-missing-imports = [
+    "tensorflow.*",
+    "tensorboard.*",
+    "absl.*",
+    "jax.*",
+    "rich.*",
+    "jaxlib.cuda.*",
+    "jaxlib.cpu.*",
+    "msgpack",
+    "numpy.*",
+    "optax.*",
+    "orbax.*",
+    "opt_einsum.*",
+    "scipy.*",
+    "libtpu.*",
+    "jaxlib.mlir.*",
+    "yaml",
+]
+
+[errors]
+missing-attribute = "ignore"

tests/run_all_tests.sh

@@ -5,6 +5,7 @@ PYTEST_OPTS=
 RUN_DOCTEST=false
 RUN_MYPY=false
 RUN_PYTEST=false
+RUN_PYREFLY=false
 RUN_PYTYPE=false
 GH_VENV=false
 
@@ -30,6 +31,9 @@ case $flag in
   --only-mypy)
   RUN_MYPY=true
   ;;
+  --only-pyrefly)
+  RUN_PYREFLY=true
+  ;;
   --use-venv)
   GH_VENV=true
   ;;
@@ -40,12 +44,13 @@ case $flag in
 esac
 done
 
-# if neither --only-doctest, --only-pytest, --only-pytype, --only-mypy is set, run all tests
-if ! $RUN_DOCTEST && ! $RUN_PYTEST && ! $RUN_PYTYPE && ! $RUN_MYPY; then
+# if neither --only-doctest, --only-pytest, --only-pytype, --only-mypy, --only-pyrefly is set, run all tests
+if ! $RUN_DOCTEST && ! $RUN_PYTEST && ! $RUN_PYTYPE && ! $RUN_MYPY && ! $RUN_PYREFLY; then
   RUN_DOCTEST=true
   RUN_PYTEST=true
   RUN_PYTYPE=true
   RUN_MYPY=true
+  RUN_PYREFLY=true
 fi
 
 # Activate cached virtual env for github CI
@@ -58,6 +63,7 @@ echo "PYTEST_OPTS: $PYTEST_OPTS"
 echo "RUN_DOCTEST: $RUN_DOCTEST"
 echo "RUN_PYTEST: $RUN_PYTEST"
 echo "RUN_MYPY: $RUN_MYPY"
+echo "RUN_PYREFLY: $RUN_PYREFLY"
 echo "RUN_PYTYPE: $RUN_PYTYPE"
 echo "GH_VENV: $GH_VENV"
 echo "WHICH PYTHON: $(which python)"
@@ -155,5 +161,11 @@ if $RUN_MYPY; then
   mypy --config pyproject.toml flax/ --show-error-codes
 fi
 
+if $RUN_PYREFLY; then
+  echo "=== RUNNING PYREFLY ==="
+  # Type-check using pyrefly.toml (currently scoped to flax/linen/linear.py).
+  pyrefly check
+fi
+
 # Return error code 0 if no real failures happened.
 echo "finished all tests."