Merge pull request #657 from google-deepmind:pytype-3.11

PiperOrigin-RevId: 588066525

test.sh

@@ -64,13 +64,8 @@ pip wheel --verbose --no-deps --no-clean dist/optax*.tar.gz
 pip install optax*.whl
 
 # Check types with pytype.
-# Note: pytype does not support 3.11 as of 25.06.23
-# See https://github.com/google/pytype/issues/1308
-if [ `python -c 'import sys; print(sys.version_info.minor)'` -lt 11 ];
-then
-  pip install pytype
-  pytype `find optax/_src/ examples -name '*.py' | xargs` -k -d import-error
-fi;
+pip install pytype
+pytype `find optax/_src/ examples -name '*.py' | xargs` -k -d import-error
 
 # Run tests using pytest.
 # Change directory to avoid importing the package from repo root.