Allow free-threaded parallel sampling

.github/workflows/tests.yml

@@ -345,6 +345,57 @@ jobs:
           name: Alternative backends py${{ matrix.python-version }} linker=${{ matrix.linker }}
           fail_ci_if_error: false
 
+  free_threading:
+    needs: changes
+    if: ${{ needs.changes.outputs.changes == 'true' }}
+    strategy:
+      matrix:
+        # A small, sampling-focused subset that exercises the thread-based parallel
+        # sampling path used on free-threaded (no-GIL) builds. The GIL-stays-disabled
+        # check is a pytest_sessionfinish hook in tests/conftest.py.
+        test-subset:
+          - |
+            tests/sampling/test_parallel.py
+            tests/step_methods/test_compound.py
+      fail-fast: false
+    runs-on: ubuntu-latest
+    env:
+      TEST_SUBSET: ${{ matrix.test-subset }}
+      # Disable the C backend (cxx=""). Its single-phase-init C extensions (e.g.
+      # lazylinker_ext) re-enable the GIL on a free-threaded build; with no C
+      # compiler PyTensor uses the free-threading-safe numba/Python backend.
+      PYTENSOR_FLAGS: cxx=
+      # Intentionally leave PYTHON_GIL unset. Setting it to 0 would force the GIL
+      # off and hide a dependency that fails to support free-threading; leaving it
+      # unset lets such a dependency re-enable the GIL, which the conftest
+      # pytest_sessionfinish hook asserts against.
+    defaults:
+      run:
+        shell: bash -leo pipefail {0}
+    steps:
+      - uses: actions/checkout@08c6903cd8c0fde910a37f88322edcfb5dd907a8 # v5.0.0
+        with:
+          persist-credentials: false
+      - uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
+        with:
+          python-version: "3.14t"
+      - name: Install pymc and test dependencies
+        run: |
+          python -m pip install --upgrade pip
+          # Install from PyPI (free-threaded wheels), not conda-forge. numba is
+          # included as the free-threading-safe PyTensor backend.
+          pip install -e . numba pytest pytest-cov
+          # Free-threaded sampling needs the (still unreleased) PyTensor
+          # free-threading support (pymc-devs/pytensor#2286); track its branch
+          # until it is released.
+          pip install "pytensor @ git+https://github.com/ricardoV94/pytensor@free_threaded"
+          python --version
+          python -c "import sys, sysconfig; print('Py_GIL_DISABLED =', sysconfig.get_config_var('Py_GIL_DISABLED'), '| gil_enabled =', sys._is_gil_enabled())"
+          pip list
+      - name: Run tests
+        run: |
+          python -m pytest -vv --durations=25 $TEST_SUBSET
+
   float32:
     needs: changes
     if: ${{ needs.changes.outputs.changes == 'true' }}
@@ -395,13 +446,14 @@ jobs:
   all_tests:
     if: ${{ always() }}
     runs-on: ubuntu-latest
-    needs: [changes, ubuntu, windows, macos, alternative_backends, float32]
+    needs: [changes, ubuntu, windows, macos, alternative_backends, free_threading, float32]
     steps:
       - name: Check build matrix status
         if: ${{ needs.changes.outputs.changes == 'true' &&
           ( needs.ubuntu.result != 'success' ||
           needs.windows.result != 'success' ||
           needs.macos.result != 'success' ||
           needs.alternative_backends.result != 'success' ||
+          needs.free_threading.result != 'success' ||
           needs.float32.result != 'success' ) }}
         run: exit 1

conda-envs/environment-alternative-backends.yml

@@ -13,13 +13,11 @@ 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.1.2,<3.2
@@ -39,7 +37,15 @@ 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.5

pymc/model/core.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 from __future__ import annotations
 
+import copy
 import functools
 import sys
 import threading
@@ -275,6 +276,26 @@ def get_extra_values(self):
 
         return {var.name: self._extra_vars_shared[var.name].get_value() for var in self._extra_vars}
 
+    def fork(self) -> ValueGradFunction:
+        """Return an independent copy for concurrent (multi-thread) use.
+
+        The copy shares the compiled code and any read-only inputs (e.g. observed
+        data) with the original, but gets its own input/output storage and its own
+        ``extra_vars`` shared variables. This makes it safe to call and to
+        ``set_extra_values`` on from a different thread than the original.
+        """
+        new = copy.copy(self)
+        swap = {}
+        new._extra_vars_shared = {}
+        for name, shared in self._extra_vars_shared.items():
+            value = shared.get_value(borrow=False)
+            fresh = pytensor.shared(value, shared.name, shape=value.shape)
+            new._extra_vars_shared[name] = fresh
+            swap[shared] = fresh
+        new._pytensor_function = self._pytensor_function.copy(share_memory=False, swap=swap or None)
+        new._extra_are_set = False
+        return new
+
     def __call__(self, grad_vars, *, extra_vars=None):
         if extra_vars is not None:
             self.set_extra_values(extra_vars)

pymc/progress_bar/rich_progress.py

@@ -149,11 +149,10 @@ def __rich_console__(self, console: Console, options: ConsoleOptions) -> RenderR
 class CustomBarColumn(BarColumn):
     """Bar column that recolors on divergences and renders a separator marker."""
 
-    def __init__(self, *args, failing_color: str = "red", **kwargs):
-        from matplotlib.colors import to_rgb
-
-        self.failing_color = failing_color
-        self.failing_rgb = [int(x * 255) for x in to_rgb(self.failing_color)]
+    def __init__(self, *args, failing_style: Style | None = None, **kwargs):
+        self.failing_style = (
+            failing_style if failing_style is not None else Style.parse("rgb(214,39,40)")
+        )
 
         super().__init__(*args, **kwargs)
 
@@ -163,8 +162,8 @@ def __init__(self, *args, failing_color: str = "red", **kwargs):
     def callbacks(self, task: Task):
         """Update bar color based on failure state."""
         if task.fields["failing"]:
-            self.complete_style = Style.parse("rgb({},{},{})".format(*self.failing_rgb))
-            self.finished_style = Style.parse("rgb({},{},{})".format(*self.failing_rgb))
+            self.complete_style = self.failing_style
+            self.finished_style = self.failing_style
         else:
             self.complete_style = self.default_complete_style
             self.finished_style = self.default_finished_style
@@ -264,7 +263,6 @@ def _create_progress_bar(
             CustomBarColumn(
                 bar_width=None,
                 table_column=Column("Progress", ratio=2),
-                failing_color="tab:red",
                 complete_style=Style.parse("rgb(31,119,180)"),
                 finished_style=Style.parse("rgb(31,119,180)"),
             ),