Bump numba upper bound to <=0.66.0 (main)

.github/workflows/numba-upper-bound.yml

@@ -1,8 +1,10 @@
 name: Bump numba upper bound
 
-# Watches PyPI for new numba releases. When one appears above our current
-# upper bound in pyproject.toml, opens a PR raising the cap so our normal
-# CI (Tests workflow) validates the new version before we ship it.
+# Watches conda-forge for new numba releases. When one appears above our
+# current upper bound in pyproject.toml, opens a PR raising the cap so our
+# normal CI (Tests workflow) validates the new version before we ship it.
+# CI installs numba from conda-forge, so gating here (rather than on PyPI)
+# ensures the bumped version is actually installable when CI runs.
 
 on:
   schedule:
@@ -31,7 +33,7 @@ jobs:
       - name: Install packaging
         run: python -m pip install --upgrade packaging
 
-      - name: Check PyPI and rewrite pin if needed
+      - name: Check conda-forge and rewrite pin if needed
         id: bump
         run: python scripts/bump_numba_upper_bound.py
 
@@ -62,7 +64,7 @@ jobs:
           remote="https://x-access-token:${GH_TOKEN}@github.com/${GITHUB_REPOSITORY}.git"
           git push "$remote" "$branch"
 
-          body="numba ${LATEST} is available on PyPI. This PR raises the upper bound in \`pyproject.toml\` to \`<=${NEW_CAP}\` so the normal Tests workflow validates the new release. Merge after CI passes."
+          body="numba ${LATEST} is available on conda-forge. This PR raises the upper bound in \`pyproject.toml\` to \`<=${NEW_CAP}\` so the normal Tests workflow validates the new release. Merge after CI passes."
           body="${body}"$'\n\n'"**Note:** GitHub suppresses workflow runs triggered by the default \`GITHUB_TOKEN\`, so CI (Tests, mypy, etc.) will not start automatically on this PR. Close and reopen it (or push an empty commit) to trigger the checks."
 
           gh pr create \

.github/workflows/test.yml

@@ -162,10 +162,15 @@ jobs:
         shell: micromamba-shell {0}
         run: |
 
+          # Install the exact numba cap from pyproject.toml so CI validates the
+          # version the bump-numba workflow pinned (fails loudly if conda-forge
+          # doesn't have it yet, rather than silently testing an older release).
+          NUMBA_VERSION=$(sed -n 's/.*"numba>=[0-9.]*,<=\([0-9.]*\)".*/\1/p' pyproject.toml | head -1)
+          echo "Installing numba==${NUMBA_VERSION}"
           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}" numpy "scipy<1.17.0" "numba==${NUMBA_VERSION}" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx libblas=*=*accelerate 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}" numpy "scipy<1.17.0" "numba==${NUMBA_VERSION}" pip graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock pytest-sphinx mkl mkl-service rich;
           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
@@ -233,7 +238,9 @@ jobs:
       - name: Install dependencies
         shell: micromamba-shell {0}
         run: |
-          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark
+          NUMBA_VERSION=$(sed -n 's/.*"numba>=[0-9.]*,<=\([0-9.]*\)".*/\1/p' pyproject.toml | head -1)
+          echo "Installing numba==${NUMBA_VERSION}"
+          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" mkl numpy scipy pip mkl-service cython pytest "numba==${NUMBA_VERSION}" jax jaxlib pytest-benchmark
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'

pyproject.toml

@@ -49,7 +49,7 @@ dependencies = [
     "setuptools>=59.0.0",
     "scipy>=1,<2",
     "numpy>=2.0",
-    "numba>=0.58,<=0.65.1",
+    "numba>=0.58,<=0.66.0",
     "filelock>=3.15",
     "etuples",
     "logical-unification",
@@ -80,7 +80,7 @@ tests = [
 ]
 rtd = ["sphinx>=5.1.0,<6", "pygments", "pydot"]
 jax = ["jax", "jaxlib"]
-numba = ["numba>=0.58,<=0.65.1", "llvmlite"]
+numba = ["numba>=0.58,<=0.66.0", "llvmlite"]
 
 [tool.setuptools.packages.find]
 include = ["pytensor*"]

scripts/bump_numba_upper_bound.py

@@ -1,9 +1,12 @@
 #!/usr/bin/env python
 """Bump the numba upper bound in pyproject.toml when a new numba release is out.
 
-Queries PyPI for the latest stable numba release and compares it against the
-current ``<=`` cap pinned in pyproject.toml. The pin is rewritten to the
-canonical ``"numba>=X.Y,<=A.B.C"`` form whenever:
+Queries conda-forge for the latest stable numba release and compares it against
+the current ``<=`` cap pinned in pyproject.toml. CI installs numba from
+conda-forge, so gating on conda-forge (rather than PyPI, which it typically
+leads by days) means the PR only opens once the version is actually installable
+in CI. The pin is rewritten to the canonical ``"numba>=X.Y,<=A.B.C"`` form
+whenever:
 
 * the latest release is newer than the cap; or
 * the existing pin is in a non-canonical shape (``<`` instead of ``<=``, or
@@ -33,22 +36,20 @@
 LOWER_RE = re.compile(r">=\d+\.\d+(?:\.\d+)?")
 UPPER_RE = re.compile(r"<(=?)(\d+\.\d+(?:\.\d+)?)")
 DEFAULT_LOWER = ">=0.58"
-PYPI_URL = "https://pypi.org/pypi/numba/json"
+CONDA_FORGE_URL = "https://api.anaconda.org/package/conda-forge/numba"
 
 
 def fetch_latest_numba() -> Version:
-    with urllib.request.urlopen(PYPI_URL, timeout=30) as resp:
+    with urllib.request.urlopen(CONDA_FORGE_URL, timeout=30) as resp:
         data = json.load(resp)
     candidates: list[Version] = []
-    for raw, files in data["releases"].items():
-        if not files:
-            continue
+    for raw in data["versions"]:
         v = Version(raw)
         if v.is_prerelease or v.is_devrelease:
             continue
         candidates.append(v)
     if not candidates:
-        raise RuntimeError("no stable numba release found on PyPI")
+        raise RuntimeError("no stable numba release found on conda-forge")
     return max(candidates)