From 37af4ecd1cde2bffc56832edac4f476d534e6557 Mon Sep 17 00:00:00 2001
From: Meesum Qazalbash <meesumqazalbash@gmail.com>
Date: Sat, 20 Dec 2025 13:21:01 +0500
Subject: [PATCH 1/3] feat(gh-2110): setup for `astral-sh/ty`

---
 Makefile       |  1 +
 pyproject.toml |  1 -
 setup.py       |  1 +
 ty.toml        | 22 ++++++++++++++++++++++
 4 files changed, 24 insertions(+), 1 deletion(-)
 create mode 100644 ty.toml

diff --git a/Makefile b/Makefile
index f9cce9207..4dae4e6b3 100644
--- a/Makefile
+++ b/Makefile
@@ -5,6 +5,7 @@ lint: FORCE
 	ruff format . --check
 	python scripts/update_headers.py --check
 	mypy --install-types --non-interactive numpyro
+	ty check -vvv --config-file ty.toml --exit-zero
 
 license: FORCE
 	python scripts/update_headers.py
diff --git a/pyproject.toml b/pyproject.toml
index 47925bf7a..7dee2c04c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -130,6 +130,5 @@ module = [
     "numpyro.primitives.*",
     "numpyro.patch.*",
     "numpyro.util.*",
-    "numpyro.distributions.transforms",
 ]
 ignore_errors = false
diff --git a/setup.py b/setup.py
index 02d951b1a..33b3eea55 100644
--- a/setup.py
+++ b/setup.py
@@ -57,6 +57,7 @@
             "pyro-api>=0.1.1",
             "scikit-learn",
             "scipy>=1.9",
+            "ty>=0.0.4",
         ],
         "dev": [
             "dm-haiku>=0.0.14; python_version >= '3.10'",
diff --git a/ty.toml b/ty.toml
new file mode 100644
index 000000000..fbd8e2264
--- /dev/null
+++ b/ty.toml
@@ -0,0 +1,22 @@
+[environment]
+python-version = "3.12"
+
+[rules]
+unresolved-import = "ignore"
+
+[src]
+include = ["numpyro"]
+exclude = [
+    "numpyro.contrib.control_flow.*",
+    "numpyro.contrib.funsor.*",
+    "numpyro.contrib.hsgp.*",
+    "numpyro.contrib.stochastic_support.*",
+    "numpyro.diagnostics.*",
+    "numpyro.handlers.*",
+    "numpyro.infer.elbo.*",
+    "numpyro.optim.*",
+    "numpyro.primitives.*",
+    "numpyro.patch.*",
+    "numpyro.util.*",
+]
+respect-ignore-files = true
\ No newline at end of file

From 86729f5a699475d5ab1dc29e8b412d143667b348 Mon Sep 17 00:00:00 2001
From: Meesum Qazalbash <meesumqazalbash@gmail.com>
Date: Sat, 20 Dec 2025 13:22:51 +0500
Subject: [PATCH 2/3] chore(gh-2110): add ty check in pre-commit hooks

---
 .pre-commit-config.yaml | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 6e17fdce9..2c3be4e01 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -20,6 +20,14 @@ repos:
         entry: mypy --install-types --non-interactive
         files: ^numpyro/
 
+      - id: ty
+        name: ty check
+        language: python
+        entry: ty
+        args: [check, --config-file, ty.toml]
+        pass_filenames: false
+        always_run: true
+        additional_dependencies: [ty]
 
   - repo: https://github.com/pre-commit/pre-commit-hooks
     rev: v6.0.0
@@ -30,9 +38,9 @@ repos:
         exclude: notebooks/
 
   - repo: https://github.com/codespell-project/codespell
-    rev: v2.3.0
+    rev: v2.4.1
     hooks:
       - id: codespell
         stages: [pre-commit, commit-msg]
-        args:
-          [--ignore-words-list, "Teh,aas,ans,dout", --check-filenames, --skip, "*.ipynb"]
+        args: [--ignore-words-list, "Teh,aas,ans,dout", --check-filenames, --skip, "*.ipynb"]
+

From 836eab9b6f817b60608990c2dd5f2dec40123f1e Mon Sep 17 00:00:00 2001
From: Meesum Qazalbash <meesumqazalbash@gmail.com>
Date: Sat, 20 Dec 2025 20:22:26 +0500
Subject: [PATCH 3/3] revert: mypy module overrides

---
 pyproject.toml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pyproject.toml b/pyproject.toml
index 7dee2c04c..47925bf7a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -130,5 +130,6 @@ module = [
     "numpyro.primitives.*",
     "numpyro.patch.*",
     "numpyro.util.*",
+    "numpyro.distributions.transforms",
 ]
 ignore_errors = false