Make mypy happy

Author: jessegrabowski

pymc/sampling/jax.py

@@ -68,6 +68,8 @@
     "sample_numpyro_nuts",
 )
 
+JaxNutsSampler = Literal["numpyro", "blackjax"]
+
 
 @jax_funcify.register(Assert)
 @jax_funcify.register(CheckParameterValue)
@@ -486,7 +488,7 @@ def sample_jax_nuts(
     postprocessing_chunks=None,
     idata_kwargs: dict | None = None,
     compute_convergence_checks: bool = True,
-    nuts_sampler: Literal["numpyro", "blackjax"],
+    nuts_sampler: JaxNutsSampler,
 ) -> az.InferenceData:
     """
     Draw samples from the posterior using a jax NUTS method.

pymc/sampling/mcmc.py

@@ -82,10 +82,11 @@
 
 Step: TypeAlias = BlockedStep | CompoundStep
 
-ExternalNutsSampler = ["nutpie", "numpyro", "blackjax"]
+ExternalNutsSampler = Literal["nutpie", "numpyro", "blackjax"]
 NutsSampler = Literal["pymc"] | ExternalNutsSampler
-
 NutpieBackend = Literal["numba", "jax"]
+
+
 NUTPIE_BACKENDS = get_args(NutpieBackend)
 NUTPIE_DEFAULT_BACKEND = cast(NutpieBackend, "numba")
 
@@ -381,6 +382,10 @@ def extract_backend(string: str) -> NutpieBackend:
     elif sampler in ("numpyro", "blackjax"):
         import pymc.sampling.jax as pymc_jax
 
+        from pymc.sampling.jax import JaxNutsSampler
+
+        sampler = cast(JaxNutsSampler, sampler)
+
         idata = pymc_jax.sample_jax_nuts(
             draws=draws,
             tune=tune,