Pin ty and improve type hints

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
         files: "(.py$)|(.*.ipynb$)"
 
   - repo: https://github.com/astral-sh/ty-pre-commit
-    rev: v0.0.49
+    rev: v0.0.57
     hooks:
       - id: ty
 

numpyro/contrib/control_flow/scan.py

@@ -69,10 +69,10 @@ def _subs_wrapper(subs_map, i, length, site):
                 # and generate a new sample otherwise
                 return lax.cond(
                     i < shape[0],
+                    lambda val, _: val[0][val[1]],
+                    lambda _, key: site["fn"](rng_key=key, sample_shape=sample_shape),
                     (value, i),
-                    lambda val: val[0][val[1]],
                     rng_key,
-                    lambda val: site["fn"](rng_key=val, sample_shape=sample_shape),
                 )
             else:
                 raise RuntimeError(

numpyro/contrib/stochastic_support/dcc.py

@@ -245,10 +245,10 @@ def _combine_inferences(
 
         def log_weight(
             rng_key: jax.Array,
-            i: int,
+            i: int | jax.Array,
             slp_model: Callable,
             slp_samples: dict[str, Any],
-        ) -> float:
+        ) -> jax.Array:
             trace = {k: v[i] for k, v in slp_samples.items()}
             guide = AutoNormal(
                 slp_model,

pyproject.toml

@@ -84,7 +84,7 @@ test = [
     "ruff>=0.1.8",
     "scikit-learn",
     "scipy>=1.9",
-    "ty>=0.0.49",
+    "ty==0.0.57",
 ]
 
 [build-system]