resolved conflict in test_sampling.py

Author: purna135

.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
@@ -128,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))
 - ...
 
 

pymc/__init__.py

@@ -13,7 +13,7 @@
 #   limitations under the License.
 
 # pylint: disable=wildcard-import
-__version__ = "4.0.0b3"
+__version__ = "4.0.0b4"
 
 import logging
 import multiprocessing as mp

pymc/aesaraf.py

@@ -941,7 +941,9 @@ def local_check_parameter_to_ninf_switch(fgraph, node):
 )
 
 
-def compile_pymc(inputs, outputs, mode=None, **kwargs):
+def compile_pymc(
+    inputs, outputs, mode=None, **kwargs
+) -> Callable[..., Union[np.ndarray, List[np.ndarray]]]:
     """Use ``aesara.function`` with specialized pymc rewrites always enabled.
 
     Included rewrites

pymc/backends/arviz.py

@@ -7,6 +7,7 @@
     Any,
     Dict,
     Iterable,
+    Mapping,
     Optional,
     Tuple,
     Union,
@@ -532,8 +533,8 @@ def to_inference_data(self):
 def to_inference_data(
     trace: Optional["MultiTrace"] = None,
     *,
-    prior: Optional[Dict[str, Any]] = None,
-    posterior_predictive: Optional[Dict[str, Any]] = None,
+    prior: Optional[Mapping[str, Any]] = None,
+    posterior_predictive: Optional[Mapping[str, Any]] = None,
     log_likelihood: Union[bool, Iterable[str]] = True,
     coords: Optional[CoordSpec] = None,
     dims: Optional[DimSpec] = None,