Address drastic slowdown in mypy runtime

.github/workflows/ci-build.yaml

@@ -23,7 +23,7 @@ permissions:
 jobs:
   lint_and_typecheck:
     runs-on: ubuntu-latest
-    timeout-minutes: 20
+    timeout-minutes: 5
     steps:
       - name: Cancel previous
         uses: styfle/cancel-workflow-action@0.11.0

jax/_src/third_party/scipy/betaln.py

@@ -3,7 +3,10 @@ import jax.numpy as jnp
 from jax._src.typing import Array, ArrayLike
 from jax._src.numpy.lax_numpy import _promote_args_inexact
 
-def algdiv(a: ArrayLike, b: ArrayLike) -> Array:
+# Note: for mysterious reasons, annotating this leads to very slow mypy runs.
+# def algdiv(a: ArrayLike, b: ArrayLike) -> Array:
+
+def algdiv(a, b):
     """
     Compute ``log(gamma(a))/log(gamma(a + b))`` when ``b >= 8``.