mypy

Author: awf

pyproject.toml

@@ -33,3 +33,8 @@ optional-dependencies = {dev = {file = ["requirements-dev.txt"]}}
 [tool.black]
 line-length = 88
 fast = true
+
+[tool.mypy]
+[[tool.mypy.overrides]]
+module = "mx.*"
+ignore_missing_imports = true

requirements-dev.txt

@@ -1,6 +1,7 @@
 # Requirements for tests
 pytest
 ml_dtypes
+microxcaling
 
 # Requirements for development
 pre-commit

test/test_microxcaling.py

@@ -34,7 +34,12 @@
     ],
     ids=str,
 )
-def test_mx(mx_round, gf_round, mx_etype, gf_etype):
+def test_mx(
+    mx_round: str,
+    gf_round: RoundMode,
+    mx_etype: ElemFormat,
+    gf_etype: FormatInfo,
+) -> None:
     A = torch.arange(32) / 2 - 5
 
     mx_specs = dict(