Add CI to run ty on jax/_src/dtypes.py

Author: jakevdp

.github/workflows/ci-build.yaml

@@ -189,3 +189,16 @@ jobs:
         JAX_PLATFORM_NAME: cpu
     - name: Run GPU tests
       run: python -m pytest examples/ffi/tests
+
+  ty_typecheck:
+    name: TypeCheck with ty
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@1af3b93b6815bc44a9784bd300feb67ff0d1eeb3  # v6.0.0
+    - uses: actions/setup-python@83679a892e2d95755f2dac6acb0bfd1e9ac5d548  # v6.1.0
+      with:
+        python-version: '3.14'
+    - run: |
+        pip install uv~=0.5.30
+        uv pip install --system ty
+    - run: uv run ty check .

jax/_src/dtypes.py

@@ -22,6 +22,7 @@
 from __future__ import annotations
 
 import abc
+import builtins
 import dataclasses
 import functools
 import types
@@ -86,7 +87,7 @@ class ExtendedDType(StrictABC):
   """Abstract Base Class for extended dtypes"""
   @property
   @abc.abstractmethod
-  def type(self) -> type: ...
+  def type(self) -> builtins.type: ...
 
 
 # fp8 support
@@ -634,7 +635,7 @@ def isdtype(dtype: DTypeLike, kind: str | DTypeLike | tuple[str | DTypeLike, ...
     True or False
   """
   the_dtype = np.dtype(dtype)
-  kind_tuple: tuple[str | DTypeLike, ...] = (
+  kind_tuple: tuple[str | DTypeLike, ...] = (  # ty: ignore[invalid-assignment]
     kind if isinstance(kind, tuple) else (kind,)
   )
   options: set[DType] = set()
@@ -914,7 +915,7 @@ def check_valid_dtype(dtype: DType) -> None:
     raise TypeError(f"Dtype {dtype} is not a valid JAX array "
                     "type. Only arrays of numeric types are supported by JAX.")
 
-def _maybe_canonicalize_explicit_dtype(dtype: DType, fun_name: str) -> DType:
+def _maybe_canonicalize_explicit_dtype(dtype: DType, fun_name: str | None) -> DType:
   "Canonicalizes explicitly requested dtypes, per explicit_x64_dtypes."
   allow = config.explicit_x64_dtypes.value
   if allow == config.ExplicitX64Mode.ALLOW or config.enable_x64.value:

pyproject.toml

@@ -159,3 +159,6 @@ max-complexity = 18
 "docs/notebooks/Custom_derivative_rules_for_Python_code.ipynb" = ["F811"]
 "docs/jep/9407-type-promotion.ipynb" = ["F811"]
 "docs/autodidax.ipynb" = ["F811"]
+
+[tool.ty.src]
+include = ["jax/_src/dtypes.py"]