Merge remote-tracking branch 'upstream/main' into feature/zoib

Author: Noah Shin

.github/workflows/tests.yml

@@ -153,6 +153,11 @@ jobs:
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
       PYTENSOR_FLAGS: linker=${{ matrix.linker }}
+      # AVX-512 codegen crashes LLVM's loop vectorizer (llvm-project#210601,
+      # fixed after the 20.1.8 llvmlite bundles). Not x86-64-v3: it fails
+      # test_sampling_consistency.
+      NUMBA_CPU_NAME: generic
+      NUMBA_CPU_FEATURES: ""
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -202,6 +207,9 @@ jobs:
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
       PYTENSOR_FLAGS: linker=${{ matrix.linker }}
+      # See the ubuntu job for why the numba codegen target is pinned.
+      NUMBA_CPU_NAME: generic
+      NUMBA_CPU_FEATURES: ""
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -311,6 +319,9 @@ jobs:
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
       PYTENSOR_FLAGS: linker=${{ matrix.linker }}
+      # See the ubuntu job for why the numba codegen target is pinned.
+      NUMBA_CPU_NAME: generic
+      NUMBA_CPU_FEATURES: ""
     defaults:
       run:
         shell: bash -leo pipefail {0}
@@ -360,6 +371,9 @@ jobs:
     env:
       TEST_SUBSET: ${{ matrix.test-subset }}
       PYTENSOR_FLAGS: floatX=float32
+      # See the ubuntu job for why the numba codegen target is pinned.
+      NUMBA_CPU_NAME: generic
+      NUMBA_CPU_FEATURES: ""
     defaults:
       run:
         shell: bash -leo pipefail {0}

conda-envs/environment-alternative-backends.yml

@@ -13,16 +13,14 @@ dependencies:
 - numba
 # nutpie is installed from git in the CI workflow (see tests.yml) so we can track
 # the arviz 1.0 -compatible branch until an upstream release ships.
-# Jaxlib version must not be greater than jax version!
-- jax>=0.4.28
-- jaxlib>=0.4.28
+# The JAX stack (jax, jaxlib, numpyro, blackjax) is installed from PyPI in the pip
+# section below rather than conda-forge -- see the note there.
 - libblas=*=*mkl
 - mkl-service
 - numpy>=1.25.0
-- numpyro>=0.8.0
 - pandas>=0.24.0
 - pip
-- pytensor>=3.0.7,<3.1
+- pytensor>=3.2.2,<3.3
 - python-graphviz
 - networkx
 - rich>=13.7.1
@@ -39,7 +37,18 @@ dependencies:
 - types-cachetools
 # blackjax now does not pull in fastprogress by default
 - fastprogress>=1.0.0
-- blackjax>=1.5
 - pip:
   - numdifftools>=0.9.40
   - mcbackend>=0.4.0
+  # conda-forge builds jaxlib from source with its own XLA/LLVM toolchain, and its
+  # 0.10.x CPU build segfaults in XLA codegen while compiling some graphs
+  # (DirichletMultinomial, MvStudentT draws) on the CI runners. The official PyPI
+  # wheels compile the same graphs fine, so install the whole JAX stack from PyPI.
+  # Jaxlib version must not be greater than jax version!
+  - jax>=0.4.28
+  - jaxlib>=0.4.28
+  - numpyro>=0.8.0
+  # blackjax 1.6 removed `window_adaptation(progress_bar=...)` and replaced
+  # `progress_bar.gen_scan_fn` with a context manager; pymc.sampling.jax still
+  # uses the old API.
+  - blackjax>=1.5,<1.6

conda-envs/environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=3.0.7,<3.1
+- pytensor>=3.2.2,<3.3
 - python-graphviz
 - networkx
 - scipy>=1.4.1

conda-envs/environment-docs.yml

@@ -11,7 +11,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=3.0.7,<3.1
+- pytensor>=3.2.2,<3.3
 - python-graphviz
 - rich>=13.7.1
 - scipy>=1.4.1

conda-envs/environment-test.yml

@@ -15,7 +15,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=3.0.7,<3.1
+- pytensor>=3.2.2,<3.3
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-dev.yml

@@ -12,7 +12,7 @@ dependencies:
 - numpy>=1.25.0
 - pandas>=0.24.0
 - pip
-- pytensor>=3.0.7,<3.1
+- pytensor>=3.2.2,<3.3
 - python-graphviz
 - networkx
 - rich>=13.7.1

conda-envs/windows-environment-test.yml

@@ -16,7 +16,7 @@ dependencies:
 - pandas>=0.24.0
 - pip
 - polyagamma
-- pytensor>=3.0.7,<3.1
+- pytensor>=3.2.2,<3.3
 - python-graphviz
 - networkx
 - rich>=13.7.1

docs/source/api/dims/distributions.rst

@@ -26,6 +26,9 @@ Scalar distributions
    Gamma
    InverseGamma
    Weibull
+   Poisson
+   NegativeBinomial
+   DiracDelta
 
 
 Vector distributions

docs/source/api/model/core.rst

@@ -5,6 +5,8 @@ Model creation and inspection
 .. autosummary::
    :toctree: generated/
 
+   BaseModel
+   FrozenModel
    Model
    modelcontext
 

docs/source/api/model/optimization.rst

@@ -5,3 +5,4 @@ Model Optimization
    :toctree: generated/
 
    freeze_dims_and_data
+   freeze_model