diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index cade4e5a5..7cc4a75a2 100644
--- a/.pre-commit-config.yaml
+++ b/.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
 
diff --git a/numpyro/contrib/control_flow/scan.py b/numpyro/contrib/control_flow/scan.py
index c0c088891..509e987a4 100644
--- a/numpyro/contrib/control_flow/scan.py
+++ b/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(
diff --git a/numpyro/contrib/stochastic_support/dcc.py b/numpyro/contrib/stochastic_support/dcc.py
index 552e17335..42f7a8926 100644
--- a/numpyro/contrib/stochastic_support/dcc.py
+++ b/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,
diff --git a/pyproject.toml b/pyproject.toml
index 59bd0ed95..f50b70bad 100644
--- a/pyproject.toml
+++ b/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]