Update arviz and various small updates

.github/uv-constraints-dev.txt

@@ -0,0 +1,2 @@
+pytensor @ git+https://github.com/pymc-devs/pytensor.git@v3
+pymc @ git+https://github.com/pymc-devs/pymc.git@v6

.github/uv-constraints-main.txt

@@ -0,0 +1 @@
+arviz < 1.0.0

.github/workflows/ci.yml

@@ -12,9 +12,12 @@ on:
 
 permissions:
   contents: read
+  pages: write
+  id-token: write
 
 jobs:
-  linux:
+  build_linux:
+    name: build linux (${{ matrix.platform.target }})
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       fail-fast: false
@@ -31,16 +34,14 @@ jobs:
       - uses: actions/setup-python@v6
         with:
           python-version: |
-            3.10
-            3.11
             3.12
             3.13
             3.14
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --interpreter 3.10 3.11 3.12 3.13 3.14 --zig
+          args: --release --out dist --interpreter 3.12 3.13 3.14 --zig
           sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
           manylinux: auto
           before-script-linux: |
@@ -50,20 +51,58 @@ jobs:
         with:
           name: wheels-linux-${{ matrix.platform.target }}
           path: dist
+
+  test_linux:
+    name: test linux (${{ matrix.target }}, ${{ matrix.python-version }}, ${{ matrix.suite }})
+    runs-on: ${{ matrix.runner }}
+    needs: build_linux
+    # stan failures block release; pymc/flow are expected red against
+    # released pymc, so allow them to fail without gating.
+    continue-on-error: ${{ matrix.suite != 'stan' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        target: [x86_64, aarch64]
+        suite: [stan, pymc, flow]
+        python-version: ["3.12", "3.13", "3.14"]
+        include:
+          - target: x86_64
+            runner: ubuntu-22.04
+          - target: aarch64
+            runner: ubuntu-22.04-arm
+    steps:
+      - uses: actions/checkout@v6
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - uses: actions/download-artifact@v8
+        with:
+          name: wheels-linux-${{ matrix.target }}
+          path: dist
       - name: pytest
         shell: bash
         run: |
           set -e
           python3 -m venv .venv
           source .venv/bin/activate
-          uv pip install 'nutpie[stan]' --find-links dist --force-reinstall
           uv pip install pytest pytest-timeout pytest-arraydiff
-          pytest -m "stan and not flow" --arraydiff
-          uv pip install 'nutpie[pymc]' --find-links dist --force-reinstall
-          uv pip install jax
-          pytest -m "pymc and not flow" --arraydiff
-          uv pip install 'nutpie[all]' --find-links dist --force-reinstall
-          pytest -m flow --arraydiff
+          uv pip install nutpie --no-deps --find-links dist --no-index
+          case "${{ matrix.suite }}" in
+            stan)
+              uv pip install 'nutpie[stan]'
+              pytest -m "stan and not flow" --arraydiff
+              ;;
+            pymc)
+              uv pip install 'nutpie[pymc]' jax --constraint .github/uv-constraints-main.txt
+              pytest -m "pymc and not flow" --arraydiff
+              ;;
+            flow)
+              uv pip install 'nutpie[all]' --constraint .github/uv-constraints-main.txt
+              pytest -m flow --arraydiff
+              ;;
+          esac
 
   # pyarrow doesn't currently seem to work on musllinux
   #musllinux:
@@ -135,24 +174,18 @@ jobs:
   #          uv pip install 'nutpie[stan]' --find-links dist --force-reinstall
   #          pytest
 
-  windows:
-    runs-on: ${{ matrix.platform.runner }}
-    strategy:
-      matrix:
-        platform:
-          - runner: windows-latest
-            target: x64
+  build_windows:
+    name: build windows
+    runs-on: windows-latest
     steps:
       - uses: actions/checkout@v6
       - uses: actions/setup-python@v6
         with:
           python-version: |
-            3.10
-            3.11
             3.12
             3.13
             3.14
-          architecture: ${{ matrix.platform.target }}
+          architecture: x64
       - name: Install uv
         uses: astral-sh/setup-uv@v7
       - name: Install LLVM and Clang
@@ -161,39 +194,68 @@ jobs:
           version: "15.0"
           directory: ${{ runner.temp }}/llvm
       - name: Set up TBB
-        if: matrix.os == 'windows-latest'
         run: |
           Add-Content $env:GITHUB_PATH "$(pwd)/stan/lib/stan_math/lib/tbb"
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         env:
           LIBCLANG_PATH: ${{ runner.temp }}/llvm/lib
         with:
-          target: ${{ matrix.platform.target }}
-          args: --release --out dist --find-interpreter
+          target: x64
+          args: --release --out dist --interpreter 3.12 3.13 3.14
           sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
       - name: Upload wheels
         uses: actions/upload-artifact@v7
         with:
-          name: wheels-windows-${{ matrix.platform.target }}
+          name: wheels-windows-x64
+          path: dist
+
+  test_windows:
+    name: test windows (${{ matrix.python-version }}, ${{ matrix.suite }})
+    runs-on: windows-latest
+    needs: build_windows
+    continue-on-error: ${{ matrix.suite != 'stan' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        suite: [stan, pymc, flow]
+        python-version: ["3.12", "3.13", "3.14"]
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - uses: actions/download-artifact@v8
+        with:
+          name: wheels-windows-x64
           path: dist
       - name: pytest
-        if: ${{ !startsWith(matrix.platform.target, 'aarch64') }}
         shell: bash
         run: |
           set -e
           python3 -m venv .venv
           source .venv/Scripts/activate
-          uv pip install "nutpie[stan]" --find-links dist --force-reinstall
           uv pip install pytest pytest-timeout pytest-arraydiff
-          pytest -m "stan and not flow" --arraydiff
-          uv pip install "nutpie[pymc]" --find-links dist --force-reinstall
-          uv pip install jax
-          pytest -m "pymc and not flow" --arraydiff
-          uv pip install "nutpie[all]" --find-links dist --force-reinstall
-          pytest -m flow --arraydiff
+          uv pip install nutpie --no-deps --find-links dist --no-index
+          case "${{ matrix.suite }}" in
+            stan)
+              uv pip install "nutpie[stan]"
+              pytest -m "stan and not flow" --arraydiff
+              ;;
+            pymc)
+              uv pip install "nutpie[pymc]" jax --constraint .github/uv-constraints-main.txt
+              pytest -m "pymc and not flow" --arraydiff
+              ;;
+            flow)
+              uv pip install "nutpie[all]" --constraint .github/uv-constraints-main.txt
+              pytest -m flow --arraydiff
+              ;;
+          esac
 
-  macos:
+  build_macos:
+    name: build macos (${{ matrix.platform.target }})
     runs-on: ${{ matrix.platform.runner }}
     strategy:
       fail-fast: false
@@ -208,8 +270,6 @@ jobs:
       - uses: actions/setup-python@v6
         with:
           python-version: |
-            3.10
-            3.11
             3.12
             3.13
             3.14
@@ -222,29 +282,145 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.platform.target }}
-          args: --release --out dist --find-interpreter
+          args: --release --out dist --interpreter 3.12 3.13 3.14 --zig
           sccache: ${{ !startsWith(github.ref, 'refs/tags/') }}
       - name: Upload wheels
         uses: actions/upload-artifact@v7
         with:
           name: wheels-macos-${{ matrix.platform.target }}
           path: dist
+
+  test_macos:
+    name: test macos (${{ matrix.target }}, ${{ matrix.python-version }}, ${{ matrix.suite }})
+    runs-on: ${{ matrix.runner }}
+    needs: build_macos
+    continue-on-error: ${{ matrix.suite != 'stan' }}
+    strategy:
+      fail-fast: false
+      matrix:
+        target: [x86_64, aarch64]
+        suite: [stan, pymc, flow]
+        python-version: ["3.12", "3.13", "3.14"]
+        include:
+          - target: x86_64
+            runner: macos-15-intel
+          - target: aarch64
+            runner: macos-14
+        exclude:
+          # numba isn't available on macOS x86_64
+          - target: x86_64
+            suite: pymc
+          - target: x86_64
+            suite: flow
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - uses: actions/download-artifact@v8
+        with:
+          name: wheels-macos-${{ matrix.target }}
+          path: dist
       - name: pytest
         run: |
           set -e
           python3 -m venv .venv
           source .venv/bin/activate
-          uv pip install 'nutpie[stan]' --find-links dist --force-reinstall
           uv pip install pytest pytest-timeout pytest-arraydiff
-          pytest -m "stan and not flow" --arraydiff
-          # Skip on x86_64 due to lack of numba support
-          if [ "${{ matrix.platform.target }}" != "x86_64" ]; then
-            uv pip install 'nutpie[pymc]' --find-links dist --force-reinstall
-            uv pip install jax
-            pytest -m "pymc and not flow" --arraydiff
-            uv pip install 'nutpie[all]' --find-links dist --force-reinstall
-            pytest -m flow --arraydiff
-          fi
+          uv pip install nutpie --no-deps --find-links dist --no-index
+          case "${{ matrix.suite }}" in
+            stan)
+              uv pip install 'nutpie[stan]'
+              pytest -m "stan and not flow" --arraydiff
+              ;;
+            pymc)
+              uv pip install 'nutpie[pymc]' jax --constraint .github/uv-constraints-main.txt
+              pytest -m "pymc and not flow" --arraydiff
+              ;;
+            flow)
+              uv pip install 'nutpie[all]' --constraint .github/uv-constraints-main.txt
+              pytest -m flow --arraydiff
+              ;;
+          esac
+
+  test_pymc_dev:
+    name: pymc (dev)
+    runs-on: ubuntu-22.04
+    needs: build_linux
+    steps:
+      - uses: actions/checkout@v6
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - uses: actions/download-artifact@v8
+        with:
+          name: wheels-linux-x86_64
+          path: dist
+      - name: pytest
+        run: |
+          python3 -m venv .venv
+          source .venv/bin/activate
+          uv pip install nutpie --no-deps --find-links dist --no-index
+          uv pip install 'nutpie[dev]' --constraint .github/uv-constraints-dev.txt
+          pytest -m "pymc" --arraydiff
+
+  docs:
+    runs-on: ubuntu-22.04
+    steps:
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
+        with:
+          python-version: "3.12"
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+      - name: Install system dependencies
+        run: sudo apt-get update && sudo apt-get install -y clang libclang-dev
+      - name: Install Quarto
+        uses: quarto-dev/quarto-actions/setup@v2
+      - name: Build and install nutpie
+        run: |
+          python3 -m venv .venv
+          source .venv/bin/activate
+          uv pip install maturin
+          maturin build --release --out dist
+      - name: Install doc dependencies
+        run: |
+          source .venv/bin/activate
+          uv pip install nutpie --no-deps --find-links dist --no-index
+          uv pip install --constraint .github/uv-constraints-dev.txt "nutpie[docs]"
+      - name: Render docs
+        env:
+          TBB_CXX_TYPE: clang
+        run: |
+          source .venv/bin/activate
+          cd docs
+          quarto render
+      - name: Upload docs artifact
+        if: ${{ github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/') }}
+        uses: actions/upload-pages-artifact@v3
+        with:
+          path: docs/_site
+
+  deploy-docs:
+    name: Deploy Docs
+    runs-on: ubuntu-latest
+    if: ${{ github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/') }}
+    needs: [docs]
+    permissions:
+      pages: write
+      id-token: write
+    environment:
+      name: github-pages
+      url: ${{ steps.deployment.outputs.page_url }}
+    steps:
+      - name: Deploy to GitHub Pages
+        id: deployment
+        uses: actions/deploy-pages@v4
+
   sdist:
     runs-on: ubuntu-latest
     steps:
@@ -264,7 +440,18 @@ jobs:
     name: Release
     runs-on: ubuntu-latest
     if: ${{ startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch' }}
-    needs: [linux, windows, macos, sdist]
+    needs:
+      [
+        build_linux,
+        build_windows,
+        build_macos,
+        sdist,
+        test_linux,
+        test_windows,
+        test_macos,
+        test_pymc_dev,
+        docs,
+      ]
     environment:
       name: pypi
     permissions: