Merge branch 'main' into adding-blackjax-support

Author: zaxtax

pymc/data.py

@@ -25,7 +25,6 @@
 import aesara
 import aesara.tensor as at
 import numpy as np
-import pandas as pd
 
 from aesara.compile.sharedvalue import SharedVariable
 from aesara.graph.basic import Apply
@@ -472,7 +471,7 @@ def determine_coords(model, value, dims: Optional[Sequence[str]] = None) -> Dict
     coords = {}
 
     # If value is a df or a series, we interpret the index as coords:
-    if isinstance(value, (pd.Series, pd.DataFrame)):
+    if hasattr(value, "index"):
         dim_name = None
         if dims is not None:
             dim_name = dims[0]
@@ -482,7 +481,7 @@ def determine_coords(model, value, dims: Optional[Sequence[str]] = None) -> Dict
             coords[dim_name] = value.index
 
     # If value is a df, we also interpret the columns as coords:
-    if isinstance(value, pd.DataFrame):
+    if hasattr(value, "columns"):
         dim_name = None
         if dims is not None:
             dim_name = dims[1]
@@ -501,7 +500,7 @@ def determine_coords(model, value, dims: Optional[Sequence[str]] = None) -> Dict
         for size, dim in zip(value.shape, dims):
             coord = model.coords.get(dim, None)
             if coord is None:
-                coords[dim] = pd.RangeIndex(size, name=dim)
+                coords[dim] = range(size)
 
     return coords
 

pymc/distributions/discrete.py

@@ -971,12 +971,12 @@ class HyperGeometric(Discrete):
 
     Parameters
     ----------
-    N : integer
-        Total size of the population
-    k : integer
-        Number of successful individuals in the population
-    n : integer
-        Number of samples drawn from the population
+    N : tensor_like of integer
+        Total size of the population (N > 0)
+    k : tensor_like of integer
+        Number of successful individuals in the population (0 <= k <= N)
+    n : tensor_like of integer
+        Number of samples drawn from the population (0 <= n <= N)
     """
 
     rv_op = hypergeometric
@@ -1004,6 +1004,10 @@ def logp(value, good, bad, n):
         value : numeric
             Value(s) for which log-probability is calculated. If the log probabilities for multiple
             values are desired the values must be provided in a numpy array or Aesara tensor
+        good : integer, array_like or TensorVariable
+            Number of successful individuals in the population. Alias for parameter :math:`k`.
+        bad : integer, array_like or TensorVariable
+            Number of unsuccessful individuals in the population. Alias for :math:`N-k`.
 
         Returns
         -------
@@ -1042,8 +1046,14 @@ def logcdf(value, good, bad, n):
 
         Parameters
         ----------
-        value: numeric
+        value : numeric
             Value for which log CDF is calculated.
+        good : integer
+            Number of successful individuals in the population. Alias for parameter :math:`k`.
+        bad : integer
+            Number of unsuccessful individuals in the population. Alias for :math:`N-k`.
+        n : integer
+            Number of samples drawn from the population (0 <= n <= N)
 
         Returns
         -------

pymc/model.py

@@ -631,9 +631,7 @@ def logp_dlogp_function(self, grad_vars=None, tempered=False, **kwargs):
                     raise ValueError(f"Can only compute the gradient of continuous types: {var}")
 
         if tempered:
-            # TODO: Should this differ from self.datalogpt,
-            #  where the potential terms are added to the observations?
-            costs = [self.varlogpt + self.potentiallogpt, self.observedlogpt]
+            costs = [self.varlogpt, self.datalogpt]
         else:
             costs = [self.logpt()]
 

pymc/sampling.py

@@ -181,14 +181,14 @@ def assign_step_methods(model, step=None, methods=STEP_METHODS, step_kwargs=None
     Parameters
     ----------
     model : Model object
-        A fully-specified model object
-    step : step function or vector of step functions
+        A fully-specified model object.
+    step : step function or iterable of step functions, optional
         One or more step functions that have been assigned to some subset of
         the model's parameters. Defaults to ``None`` (no assigned variables).
-    methods : vector of step method classes
+    methods : iterable of step method classes, optional
         The set of step methods from which the function may choose. Defaults
         to the main step methods provided by PyMC.
-    step_kwargs : dict
+    step_kwargs : dict, optional
         Parameters for the samplers. Keys are the lower case names of
         the step method, values a dict of arguments.
 

pymc/sampling_jax.py

@@ -1,12 +1,12 @@
-# pylint: skip-file
 import os
 import re
 import sys
 import warnings
 
 from functools import partial
-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", "")
@@ -132,7 +132,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:
@@ -144,8 +144,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
@@ -154,10 +159,8 @@ def _get_batched_jittered_initial_points(
         list with one item per variable and number of chains as batch dimension.
         Each item has shape `(chains, *var.shape)`
     """
-    if isinstance(random_seed, (int, np.integer)):
-        random_seed = np.random.default_rng(random_seed).integers(2**30, size=chains)
-    elif not isinstance(random_seed, (list, tuple, np.ndarray)):
-        raise ValueError(f"The `seeds` must be int or array-like. Got {type(random_seed)} instead.")
+
+    random_seed = np.random.default_rng(random_seed).integers(2**30, size=chains)
 
     assert len(random_seed) == chains
 
@@ -373,19 +376,19 @@ def sample_blackjax_nuts(
 
 
 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,
 ):
     """
     Draw samples from the posterior using the NUTS method from the ``numpyro`` library.
@@ -456,9 +459,7 @@ def sample_numpyro_nuts(
         dims = {}
 
     if random_seed is None:
-        random_seed = model.rng_seeder.randint(
-            2**30, dtype=np.int64, size=chains if chains > 1 else None
-        )
+        random_seed = model.rng_seeder.randint(2**30, dtype=np.int64)
 
     tic1 = datetime.now()
     print("Compiling...", file=sys.stdout)