Merge branch 'main' into grw

Author: canyon289

.github/workflows/arviz_compat.yml

@@ -82,7 +82,7 @@ jobs:
         run: |
           conda activate pymc-test-py39
           pip uninstall arviz -y
-          pip install git+git://github.com/arviz-devs/arviz.git
+          pip install git+https://github.com/arviz-devs/arviz
       - name: Run tests
         run: |
           python -m pytest -vv --cov=pymc --cov-report=xml --cov-report term --durations=50 $TEST_SUBSET

.pre-commit-config.yaml

@@ -14,7 +14,7 @@ repos:
         exclude: ^requirements-dev\.txt$
     -   id: trailing-whitespace
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v0.931
+  rev: v0.940
   hooks:
     - id: mypy
       name: Run static type checks
@@ -37,7 +37,7 @@ repos:
     - id: isort
       name: isort
 - repo: https://github.com/asottile/pyupgrade
-  rev: v2.31.0
+  rev: v2.31.1
   hooks:
     - id: pyupgrade
       args: [--py37-plus]

RELEASE-NOTES.md

@@ -7,8 +7,8 @@ Instead update the vNext section until 4.0.0 is out.
 ⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠⚠
 -->
 
-## PyMC vNext (4.0.0b1 → 4.0.0b2 → 4.0.0b3 → 4.0.0b4 → 4.0.0)
-⚠ The changes below are the delta between the upcoming releases `v3.11.5` →...→ `v4.0.0`.
+## PyMC vNext (4.0.0b1 → ... → 4.0.0b4 → 4.0.0)
+⚠ The changes below are the delta between the releases `v3.11.5` →...→ `v4.0.0`.
 
 ### Not-yet working features
 We plan to get these working again, but at this point their inner workings have not been refactored.
@@ -74,6 +74,7 @@ All of the above apply to:
   - The function `replace_with_values` function has been added to `gp.utils`.
   - `MarginalSparse` has been renamed `MarginalApprox`.
   - Removed `MixtureSameFamily`. `Mixture` is now capable of handling batched multivariate components (see [#5438](https://github.com/pymc-devs/pymc/pull/5438)).
+  - `ZeroInflatedPoisson` `theta` parameter was renamed to `mu` (see [#5584](https://github.com/pymc-devs/pymc/pull/5584)).
 - ...
 
 ### Expected breaks
@@ -96,6 +97,7 @@ All of the above apply to:
 This includes API changes we did not warn about since at least `3.11.0` (2021-01).
 
 - Setting initial values through `pm.Distribution(testval=...)` is now `pm.Distribution(initval=...)`.
+- Alternative `sd` keyword argument has been removed from all distributions. `sigma` should be used instead (see [#5583](https://github.com/pymc-devs/pymc/pull/5583)).
 
 
 ### New features
@@ -127,6 +129,9 @@ This includes API changes we did not warn about since at least `3.11.0` (2021-01
 - Univariate censored distributions are now available via `pm.Censored`. [#5169](https://github.com/pymc-devs/pymc/pull/5169)
 - Nested models now inherit the parent model's coordinates. [#5344](https://github.com/pymc-devs/pymc/pull/5344)
 - `softmax` and `log_softmax` functions added to `math` module (see [#5279](https://github.com/pymc-devs/pymc/pull/5279)).
+- Adding support for blackjax's NUTS sampler `pymc.sampling_jax` (see [#5477](ihttps://github.com/pymc-devs/pymc/pull/5477))
+- `pymc.sampling_jax` samplers support `log_likelihood`, `observed_data`, and `sample_stats` in returned InferenceData object (see [#5189](https://github.com/pymc-devs/pymc/pull/5189))
+- Adding support for `pm.Deterministic` in `pymc.sampling_jax` (see [#5182](https://github.com/pymc-devs/pymc/pull/5182))
 - ...
 
 

benchmarks/benchmarks/benchmarks.py

@@ -32,17 +32,17 @@ def glm_hierarchical_model(random_seed=123):
 
     n_counties = len(data.county.unique())
     with pm.Model() as model:
-        mu_a = pm.Normal("mu_a", mu=0.0, sd=100**2)
+        mu_a = pm.Normal("mu_a", mu=0.0, sigma=100**2)
         sigma_a = pm.HalfCauchy("sigma_a", 5)
-        mu_b = pm.Normal("mu_b", mu=0.0, sd=100**2)
+        mu_b = pm.Normal("mu_b", mu=0.0, sigma=100**2)
         sigma_b = pm.HalfCauchy("sigma_b", 5)
-        a = pm.Normal("a", mu=0, sd=1, shape=n_counties)
-        b = pm.Normal("b", mu=0, sd=1, shape=n_counties)
+        a = pm.Normal("a", mu=0, sigma=1, shape=n_counties)
+        b = pm.Normal("b", mu=0, sigma=1, shape=n_counties)
         a = mu_a + sigma_a * a
         b = mu_b + sigma_b * b
         eps = pm.HalfCauchy("eps", 5)
         radon_est = a[county_idx] + b[county_idx] * data.floor.values
-        pm.Normal("radon_like", mu=radon_est, sd=eps, observed=data.log_radon)
+        pm.Normal("radon_like", mu=radon_est, sigma=eps, observed=data.log_radon)
     return model
 
 
@@ -58,7 +58,7 @@ def mixture_model(random_seed=1234):
 
     with pm.Model() as model:
         w = pm.Dirichlet("w", a=np.ones_like(w_true))
-        mu = pm.Normal("mu", mu=0.0, sd=10.0, shape=w_true.shape)
+        mu = pm.Normal("mu", mu=0.0, sigma=10.0, shape=w_true.shape)
         enforce_order = pm.Potential(
             "enforce_order",
             at.switch(mu[0] - mu[1] <= 0, 0.0, -np.inf)
@@ -88,7 +88,7 @@ class OverheadSuite:
     def setup(self, step):
         self.n_steps = 10000
         with pm.Model() as self.model:
-            pm.Normal("x", mu=0, sd=1)
+            pm.Normal("x", mu=0, sigma=1)
 
     def time_overhead_sample(self, step):
         with self.model:
@@ -133,8 +133,8 @@ def time_drug_evaluation(self):
         sigma_low = 1
         sigma_high = 10
         with pm.Model():
-            group1_mean = pm.Normal("group1_mean", y_mean, sd=y_std)
-            group2_mean = pm.Normal("group2_mean", y_mean, sd=y_std)
+            group1_mean = pm.Normal("group1_mean", y_mean, sigma=y_std)
+            group2_mean = pm.Normal("group2_mean", y_mean, sigma=y_std)
             group1_std = pm.Uniform("group1_std", lower=sigma_low, upper=sigma_high)
             group2_std = pm.Uniform("group2_std", lower=sigma_low, upper=sigma_high)
             lambda_1 = group1_std**-2
@@ -301,7 +301,7 @@ def freefall(y, t, p):
             # If we know one of the parameter values, we can simply pass the value.
             ode_solution = ode_model(y0=[0], theta=[gamma, 9.8])
             # The ode_solution has a shape of (n_times, n_states)
-            Y = pm.Normal("Y", mu=ode_solution, sd=sigma, observed=y)
+            Y = pm.Normal("Y", mu=ode_solution, sigma=sigma, observed=y)
 
             t0 = time.time()
             idata = pm.sample(500, tune=1000, chains=2, cores=2, random_seed=0)

docs/source/PyMC_and_Aesara.rst

@@ -188,8 +188,8 @@ example::
 
     with pm.Model() as model:
         mu = pm.Normal('mu', 0, 1)
-        sd = pm.HalfNormal('sd', 1)
-        y = pm.Normal('y', mu=mu, sigma=sd, observed=data)
+        sigma = pm.HalfNormal('sigma', 1)
+        y = pm.Normal('y', mu=mu, sigma=sigma, observed=data)
 
 is roughly equivalent to this::
 
@@ -203,10 +203,10 @@ is roughly equivalent to this::
     model.add_free_variable(sd_log__)
     model.add_logp_term(corrected_logp_half_normal(sd_log__))
 
-    sd = at.exp(sd_log__)
-    model.add_deterministic_variable(sd)
+    sigma = at.exp(sd_log__)
+    model.add_deterministic_variable(sigma)
 
-    model.add_logp_term(pm.Normal.dist(mu, sd).logp(data))
+    model.add_logp_term(pm.Normal.dist(mu, sigma).logp(data))
 
 The return values of the variable constructors are subclasses
 of Aesara variables, so when we define a variable we can use any
@@ -217,5 +217,5 @@ Aesara operation on them::
         # beta is a at.dvector
         beta = pm.Normal('beta', 0, 1, shape=len(design_matrix))
         predict = at.dot(design_matrix, beta)
-        sd = pm.HalfCauchy('sd', beta=2.5)
-        pm.Normal('y', mu=predict, sigma=sd, observed=data)
+        sigma = pm.HalfCauchy('sigma', beta=2.5)
+        pm.Normal('y', mu=predict, sigma=sigma, observed=data)

docs/source/contributing/developer_guide.rst

@@ -888,8 +888,8 @@ others. The challenge and some summary of the solution could be found in Luciano
 
     with pm.Model() as m:
         mu = pm.Normal('mu', 0., 1., shape=(5, 1))
-        sd = pm.HalfNormal('sd', 5., shape=(1, 10))
-        pm.Normal('x', mu=mu, sigma=sd, observed=np.random.randn(2, 5, 10))
+        sigma = pm.HalfNormal('sigma', 5., shape=(1, 10))
+        pm.Normal('x', mu=mu, sigma=sigma, observed=np.random.randn(2, 5, 10))
         trace = pm.sample_prior_predictive(100)
 
     trace['x'].shape # ==> should be (100, 2, 5, 10)