Allow building in free-threaded python

.github/workflows/test.yml

@@ -93,6 +93,7 @@ jobs:
         install-torch: [0]
         install-mlx: [0]
         install-xarray: [0]
+        free-threading: [0]
         part:
           - [ "*rest", "tests --ignore=tests/scan --ignore=tests/tensor --ignore=tests/xtensor --ignore=tests/link/numba" ]
           - [ "scan", "tests/scan" ]
@@ -142,6 +143,12 @@ jobs:
             default-mode: "CVM"
             python-version: "3.14"
             os: "macos-15"
+          # conftest's pytest_sessionfinish asserts the GIL was never re-enabled.
+          - part: ["free-threading smoke test", "tests/tensor/test_math_scipy.py"]
+            default-mode: "FAST_RUN"
+            python-version: "3.14"
+            os: "ubuntu-latest"
+            free-threading: 1
 
     steps:
       - uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6.0.3
@@ -155,7 +162,12 @@ jobs:
           micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
           init-shell: bash
           post-cleanup: "all"
-          create-args: python=${{ matrix.python-version }}
+          # Create the env with python-freethreading up front so it resolves to the
+          # cp*t build from the start; adding it later conflicts with an already-solved
+          # regular CPython.
+          create-args: >-
+            python=${{ matrix.python-version }}
+            ${{ matrix.free-threading == 1 && 'python-freethreading' || '' }}
 
       - name: Create matrix id
         id: matrix-id
@@ -171,11 +183,17 @@ jobs:
         shell: micromamba-shell {0}
         run: |
 
+          # python-freethreading constrains the solve to the free-threaded CPython build
+          if [[ $FREE_THREADING == "1" ]]; then FT_PKG="python-freethreading"; else FT_PKG=""; fi
           if [[ $OS == "macos-15" ]]; then
-            micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate rich;
+            micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate rich;
+          elif [[ $FREE_THREADING == "1" ]]; then
+            # mkl/mkl-service have no free-threaded (cp*t) builds, so rely on the default (openblas) BLAS.
+            micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx rich;
           else
-            micromamba install --yes -q "python~=${PYTHON_VERSION}" numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service rich;
+            micromamba install --yes -q "python~=${PYTHON_VERSION}" $FT_PKG numpy "scipy<1.17.0" "numba>=0.63" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service rich;
           fi
+          if [[ $FREE_THREADING == "1" ]]; then python -c 'import sysconfig; assert sysconfig.get_config_var("Py_GIL_DISABLED") == 1, "Expected a free-threaded interpreter"'; fi
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" && pip install "jax>=0.8,<0.9.1" jaxlib numpyro equinox tfp-nightly; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
           if [[ $INSTALL_MLX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" "mlx<0.29.4"; fi
@@ -186,7 +204,8 @@ jobs:
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
           if [[ $OS == "macos-15" ]]; then
             python -c 'import pytensor; assert pytensor.config.blas__ldflags.startswith("-framework Accelerate"), "Blas flags are not set to MacOS Accelerate"';
-          else
+          elif [[ $FREE_THREADING != "1" ]]; then
+            # Blas flags are only used by the C backend; the free-threading job runs the numba backend.
             python -c 'import pytensor; assert pytensor.config.blas__ldflags != "", "Blas flags are empty"';
           fi
         env:
@@ -195,13 +214,17 @@ jobs:
           INSTALL_TORCH: ${{ matrix.install-torch }}
           INSTALL_XARRAY: ${{ matrix.install-xarray }}
           INSTALL_MLX: ${{ matrix.install-mlx }}
+          FREE_THREADING: ${{ matrix.free-threading }}
           OS: ${{ matrix.os}}
 
       - name: Run tests
         shell: micromamba-shell {0}
         run: |
           if [[ $DEFAULT_MODE == "FAST_COMPILE" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,mode=FAST_COMPILE; fi
           if [[ $DEFAULT_MODE == "CVM" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,linker=cvm; fi
+          # Disable the C backend on the free-threaded build so no test imports a C
+          # extension, which would re-enable the GIL (checked in conftest).
+          if [[ $FREE_THREADING == "1" ]]; then export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,cxx=; fi
           export PYTENSOR_FLAGS=$PYTENSOR_FLAGS,on_opt_error=raise,on_shape_error=raise,gcc__cxxflags=-pipe
           python -m pytest -r A --verbose --runslow --durations=50 --cov=pytensor/ --cov-report=xml:coverage/coverage-${MATRIX_ID}.xml --no-cov-on-fail $PART --benchmark-skip
         env:
@@ -211,6 +234,7 @@ jobs:
           OMP_NUM_THREADS: 1
           PART: ${{ matrix.part[1] }}
           DEFAULT_MODE: ${{ matrix.default-mode }}
+          FREE_THREADING: ${{ matrix.free-threading }}
 
       - name: Upload coverage file
         uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1

conftest.py

@@ -1,4 +1,6 @@
 import os
+import sys
+import sysconfig
 
 import pytest
 
@@ -17,6 +19,20 @@ def pytest_sessionstart(session):
     os.environ["NUMBA_BOUNDSCHECK"] = "1"
 
 
+def pytest_sessionfinish(session, exitstatus):
+    # On a free-threaded build, importing any single-phase-init C extension
+    # re-enables the GIL. The default (numba) backend must not, so if the GIL is
+    # back on, some test pulled in PyTensor's own C extensions. Checked after the
+    # whole session so it is independent of collection order.
+    if sysconfig.get_config_var("Py_GIL_DISABLED") == 1 and sys._is_gil_enabled():
+        print(  # noqa: T201
+            "\nERROR: the GIL was re-enabled during the test session. "
+            "A C-extension was imported on a free-threaded build.",
+            file=sys.stderr,
+        )
+        session.exitstatus = pytest.ExitCode.TESTS_FAILED
+
+
 def pytest_addoption(parser):
     parser.addoption(
         "--runslow", action="store_true", default=False, help="run slow tests"

pyproject.toml

@@ -33,6 +33,7 @@ classifiers = [
     "Programming Language :: Python :: 3.12",
     "Programming Language :: Python :: 3.13",
     "Programming Language :: Python :: 3.14",
+    "Programming Language :: Python :: Free Threading :: 1 - Unstable",
 ]
 
 keywords = [
@@ -182,11 +183,18 @@ files = ["pytensor", "tests"]
 build = "*"
 # Uncomment to skip builds that compile but fail when trying to test (maybe due to incompatibility with runner)
 # archs = ["auto64"]
-# Disable any-platform (pp*), 32-bit builds, and free-threaded builds (cp*t-*)
+# Disable any-platform (pp*) and 32-bit builds
 # Additional options to consider: "*musllinux*"
-skip = ["pp*", "*-win32", "*-manylinux_i686", "cp*t-*"]
+skip = ["pp*", "*-win32", "*-manylinux_i686"]
+# Free-threaded (cp*t-*) wheels are built by default. The compiled extensions
+# are not marked free-threading-safe, so importing them re-enables the GIL
+# (with a warning). That's fine: they're no longer loaded by default.
 build-frontend = "build"
-test-command = 'python -c "import pytensor; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version())"'
+# Import the compiled extensions and evaluate a function to check the wheel loads and
+# runs everywhere, including free-threaded builds. cxx="" disables the C backend (the
+# minimal wheel-test env has no libpython to link a runtime-compiled thunk against) and
+# FAST_COMPILE avoids numba, which crashes on free-threaded Windows.
+test-command = '''python -c "import pytensor; pytensor.config.cxx = ''; import pytensor.tensor as pt; print(pytensor.__version__); from pytensor.scan import scan_perform; print(scan_perform.get_version()); x = pt.vector(); f = pytensor.function([x], (x * 2).sum(), mode='FAST_COMPILE'); assert float(f([1.0, 2.0, 3.0])) == 12.0"'''
 test-skip = ["*musllinux*", "*i686*"]
 
 # Testing seems to be running into issues locating libs where expected