Add type hints to several functions in sampling_jax.py

Author: ricardoV94

pymc/sampling_jax.py

@@ -4,8 +4,9 @@
 import sys
 import warnings
 
-from typing import Callable, List, Optional
+from typing import Callable, Dict, List, Optional, Sequence, Union
 
+from pymc.initial_point import StartDict
 from pymc.sampling import _init_jitter
 
 xla_flags = os.getenv("XLA_FLAGS", "")
@@ -130,7 +131,7 @@ def _sample_stats_to_xarray(posterior):
     return data
 
 
-def _get_log_likelihood(model, samples):
+def _get_log_likelihood(model: Model, samples) -> Dict:
     """Compute log-likelihood for all observations"""
     data = {}
     for v in model.observed_RVs:
@@ -142,8 +143,13 @@ def _get_log_likelihood(model, samples):
 
 
 def _get_batched_jittered_initial_points(
-    model, chains, initvals, random_seed, jitter=True, jitter_max_retries=10
-):
+    model: Model,
+    chains: int,
+    initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]],
+    random_seed: int,
+    jitter: bool = True,
+    jitter_max_retries: int = 10,
+) -> Union[np.ndarray, List[np.ndarray]]:
     """Get jittered initial point in format expected by NumPyro MCMC kernel
 
     Returns
@@ -173,19 +179,19 @@ def _get_batched_jittered_initial_points(
 
 
 def sample_numpyro_nuts(
-    draws=1000,
-    tune=1000,
-    chains=4,
-    target_accept=0.8,
-    random_seed=None,
-    initvals=None,
-    model=None,
+    draws: int = 1000,
+    tune: int = 1000,
+    chains: int = 4,
+    target_accept: float = 0.8,
+    random_seed: int = None,
+    initvals: Optional[Union[StartDict, Sequence[Optional[StartDict]]]] = None,
+    model: Optional[Model] = None,
     var_names=None,
-    progress_bar=True,
-    keep_untransformed=False,
-    chain_method="parallel",
-    idata_kwargs=None,
-    nuts_kwargs=None,
+    progress_bar: bool = True,
+    keep_untransformed: bool = False,
+    chain_method: str = "parallel",
+    idata_kwargs: Optional[Dict] = None,
+    nuts_kwargs: Optional[Dict] = None,
 ):
     from numpyro.infer import MCMC, NUTS