Merge branch 'main' into mlx-poc

Author: cetagostini

.github/workflows/pypi.yml

@@ -3,7 +3,6 @@ on:
   push:
     branches:
       - main
-      - auto-release
   pull_request:
     branches: [main]
   release:
@@ -16,10 +15,50 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.event_name == 'pull_request' && github.head_ref || github.sha }}
   cancel-in-progress: true
 
+permissions: {}
+
 jobs:
+  check_changes:
+    runs-on: ubuntu-latest
+    outputs:
+      should_run: ${{ steps.set_should_run.outputs.should_run }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          persist-credentials: false
+      - uses: dorny/paths-filter@de90cc6fb38fc0963ad72b210f1f284cd68cea36  # v3.0.2
+        id: filter
+        with:
+          filters: |
+            any_changed:
+              - '.github/workflows/pypi.yml'
+              - 'pyproject.toml'
+              - 'setup.py'
+              - 'pytensor/_version.py'
+              - 'pytensor/scan_perform.pyx'
+              - 'pytensor/scan_perform_ext.py'
+      - name: Set should_run output
+        id: set_should_run
+        run: |
+          if [[ "${{ github.event_name == 'release' || 
+                github.ref == 'refs/heads/main' ||
+                (
+                  github.event_name == 'pull_request'
+                  && steps.filter.outputs.any_changed == 'true'
+                )
+              }}" == "true" ]]; then
+            echo "should_run=true" >> $GITHUB_OUTPUT
+          else
+            echo "should_run=false" >> $GITHUB_OUTPUT
+          fi
+
   # The job to build precompiled pypi wheels.
   make_sdist:
     name: Make SDist
+    needs: check_changes
+    # Run if it's a release or if relevant files changed on main
+    if: |
+      needs.check_changes.outputs.should_run == 'true'
     runs-on: ubuntu-latest
     permissions:
       # write id-token and attestations are required to attest build provenance
@@ -49,6 +88,10 @@ jobs:
 
   run_checks:
     name: Build & inspect our package.
+    needs: check_changes
+    # Run if it's a release or if relevant files changed on main
+    if: |
+      needs.check_changes.outputs.should_run == 'true'
     # Note: the resulting builds are not actually published.
     # This is purely for additional testing and diagnostic purposes.
     runs-on: ubuntu-latest
@@ -62,6 +105,10 @@ jobs:
 
   build_wheels:
     name: Build wheels for ${{ matrix.platform }}
+    needs: check_changes
+    # Run if it's a release or if relevant files changed on main
+    if: |
+      needs.check_changes.outputs.should_run == 'true'
     runs-on: ${{ matrix.platform }}
     permissions:
       # write id-token and attestations are required to attest build provenance
@@ -80,7 +127,7 @@ jobs:
           persist-credentials: false
 
       - name: Build wheels
-        uses: pypa/[email protected].0
+        uses: pypa/cibuildwheel@faf86a6ed7efa889faf6996aa23820831055001a # v2.23.3
 
       - name: Attest GitHub build provenance
         uses: actions/attest-build-provenance@v2
@@ -96,6 +143,10 @@ jobs:
 
   build_universal_wheel:
     name: Build universal wheel for Pyodide
+    needs: check_changes
+    # Run if it's a release or if relevant files changed on main
+    if: |
+      needs.check_changes.outputs.should_run == 'true'
     runs-on: ubuntu-latest
     permissions:
       # write id-token and attestations are required to attest build provenance
@@ -113,7 +164,7 @@ jobs:
           python-version: '3.11'
 
       - name: Install dependencies
-        run: pip install numpy versioneer wheel
+        run: pip install --upgrade setuptools numpy versioneer wheel
 
       - name: Build universal wheel
         run: |
@@ -133,7 +184,7 @@ jobs:
 
   check_dist:
     name: Check dist
-    needs: [make_sdist,build_wheels]
+    needs: [check_changes, make_sdist, build_wheels]
     runs-on: ubuntu-22.04
     steps:
       - uses: actions/download-artifact@v4

.github/workflows/test.yml

@@ -83,11 +83,12 @@ jobs:
         install-jax: [0]
         install-torch: [0]
         install-mlx: [0]
+        install-xarray: [0]
         part:
-          - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse"
+          - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse --ignore=tests/xtensor"
           - "tests/scan"
           - "tests/sparse"
-          - "tests/tensor --ignore=tests/tensor/conv --ignore=tests/tensor/rewriting --ignore=tests/tensor/test_math.py --ignore=tests/tensor/test_basic.py --ignore=tests/tensor/test_blas.py --ignore=tests/tensor/test_math_scipy.py --ignore=tests/tensor/test_inplace.py --ignore=tests/tensor/test_elemwise.py"
+          - "tests/tensor --ignore=tests/tensor/conv --ignore=tests/tensor/rewriting --ignore=tests/tensor/test_math.py --ignore=tests/tensor/test_basic.py --ignore=tests/tensor/test_inplace.py --ignore=tests/tensor/test_blas.py --ignore=tests/tensor/test_elemwise.py --ignore=tests/tensor/test_math_scipy.py"
           - "tests/tensor/conv"
           - "tests/tensor/rewriting"
           - "tests/tensor/test_math.py"
@@ -108,7 +109,7 @@ jobs:
             python-version: "3.13"
         include:
           - os: "ubuntu-latest"
-            part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
+            part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link --ignore=pytensor/ipython.py"
             python-version: "3.12"
             numpy-version: ">=2.0"
             fast-compile: 0
@@ -117,6 +118,7 @@ jobs:
             install-jax: 0
             install-torch: 0
             install-mlx: 0
+            install-xarray: 0
           - install-numba: 1
             os: "ubuntu-latest"
             python-version: "3.10"
@@ -159,6 +161,13 @@ jobs:
             fast-compile: 0
             float32: 0
             part: "tests/link/mlx"
+          - install-xarray: 1
+            os: "ubuntu-latest"
+            python-version: "3.13"
+            numpy-version: ">=2.0"
+            fast-compile: 0
+            float32: 0
+            part: "tests/xtensor"
           - os: macos-15
             python-version: "3.13"
             numpy-version: ">=2.0"
@@ -206,6 +215,7 @@ jobs:
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" jax jaxlib numpyro && pip install tensorflow-probability; 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; fi
+          if [[ $INSTALL_XARRAY == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}" xarray xarray-einstats; fi
           pip install pytest-sphinx
           
           pip install -e ./
@@ -223,6 +233,7 @@ jobs:
           INSTALL_JAX: ${{ matrix.install-jax }}
           INSTALL_TORCH: ${{ matrix.install-torch}}
           INSTALL_MLX: ${{ matrix.install-mlx }}
+          INSTALL_XARRAY: ${{ matrix.install-xarray }}
           OS: ${{ matrix.os}}
 
       - name: Run tests

doc/conf.py

@@ -1,6 +1,7 @@
 import os
 import inspect
 import sys
+
 import pytensor
 from pathlib import Path
 
@@ -12,6 +13,7 @@
     "sphinx.ext.autodoc",
     "sphinx.ext.todo",
     "sphinx.ext.doctest",
+    "sphinx_copybutton",
     "sphinx.ext.napoleon",
     "sphinx.ext.linkcode",
     "sphinx.ext.mathjax",
@@ -86,8 +88,7 @@
 
 # List of directories, relative to source directories, that shouldn't be
 # searched for source files.
-exclude_dirs = ["images", "scripts", "sandbox"]
-exclude_patterns = ['page_footer.md', '**/*.myst.md']
+exclude_patterns = ["README.md", "images/*", "page_footer.md", "**/*.myst.md"]
 
 # The reST default role (used for this markup: `text`) to use for all
 # documents.
@@ -235,24 +236,41 @@
 # Resolve function
 # This function is used to populate the (source) links in the API
 def linkcode_resolve(domain, info):
-    def find_source():
+    def find_obj() -> object:
         # try to find the file and line number, based on code from numpy:
         # https://github.com/numpy/numpy/blob/master/doc/source/conf.py#L286
         obj = sys.modules[info["module"]]
         for part in info["fullname"].split("."):
             obj = getattr(obj, part)
+        return obj
 
+    def find_source(obj):
         fn = Path(inspect.getsourcefile(obj))
-        fn = fn.relative_to(Path(__file__).parent)
+        fn = fn.relative_to(Path(pytensor.__file__).parent)
         source, lineno = inspect.getsourcelines(obj)
         return fn, lineno, lineno + len(source) - 1
 
+    def fallback_source():
+        return info["module"].replace(".", "/") + ".py"
+
     if domain != "py" or not info["module"]:
         return None
+
     try:
-        filename = "pytensor/%s#L%d-L%d" % find_source()
+        obj = find_obj()
     except Exception:
-        filename = info["module"].replace(".", "/") + ".py"
+        filename = fallback_source()
+    else:
+        try:
+            filename = "pytensor/%s#L%d-L%d" % find_source(obj)
+        except Exception:
+            # warnings.warn(f"Could not find source code for {domain}:{info}")
+            try:
+                filename = obj.__module__.replace(".", "/") + ".py"
+            except AttributeError:
+                # Some objects do not have a __module__ attribute (?)
+                filename = fallback_source()
+
     import subprocess
 
     tag = subprocess.Popen(

doc/environment.yml

@@ -13,7 +13,9 @@ dependencies:
   - mock
   - pillow
   - pymc-sphinx-theme
+  - sphinx-copybutton
   - sphinx-design
+  - sphinx-sitemap
   - pygments
   - pydot
   - ipython
@@ -23,5 +25,4 @@ dependencies:
   - ablog
   - pip
   - pip:
-    - sphinx_sitemap
     - -e ..