Allow free-threaded parallel sampling

Author: ricardoV94

.github/workflows/tests.yml

@@ -345,6 +345,53 @@ 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 }}
+      # 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 +442,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

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)