merge with master

Author: aseyboldt

.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
@@ -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
@@ -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

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 ..

doc/library/index.rst

@@ -20,14 +20,12 @@ Modules
    d3viz/index
    graph/index
    gradient
-   misc/pkl_utils
    printing
-   scalar/index
    scan
    sparse/index
-   sparse/sandbox
    tensor/index
    typed_list
+   xtensor/index
 
 .. module:: pytensor
    :platform: Unix, Windows

doc/library/xtensor/index.md

@@ -0,0 +1,101 @@
+(libdoc_xtensor)=
+# `xtensor` -- XTensor operations
+
+This module implements as abstraction layer on regular tensor operations, that behaves like Xarray.
+
+A new type {class}`pytensor.xtensor.type.XTensorType`, generalizes the {class}`pytensor.tensor.TensorType` 
+with the addition of a `dims` attribute, that labels the dimensions of the tensor. 
+
+Variables of XTensorType (i.e.,  {class}`pytensor.xtensor.type.XTensorVariable`s) are the symbolic counterpart
+to xarray DataArray objects.
+
+The module implements several PyTensor operations {class}`pytensor.xtensor.basic.XOp`s, whose signature mimics that of 
+xarray (and xarray_einstats) DataArray operations. These operations, unlike most regular PyTensor operations, cannot 
+be directly evaluated, but require a rewrite (lowering) into a regular tensor graph that can itself be evaluated as usual.
+
+Like regular PyTensor, we don't need an Op for every possible method or function in the public API of xarray.
+If the existing XOps can be composed to produce the desired result, then we can use them directly.
+
+## Coordinates
+For now, there's no analogous of xarray coordinates, so you won't be able to do coordinate operations like `.sel`.
+The graphs produced by an xarray program without coords are much more amenable to the numpy-like backend of PyTensor.
+Coords involve aspects of Pandas/database query and joining that are not trivially expressible in PyTensor.
+
+## Example
+
+
+```{testcode}
+
+import pytensor.tensor as pt
+import pytensor.xtensor as ptx
+
+a = pt.tensor("a", shape=(3,))
+b = pt.tensor("b", shape=(4,))
+
+ax = ptx.as_xtensor(a, dims=["x"])
+bx = ptx.as_xtensor(b, dims=["y"])
+
+zx = ax + bx
+assert zx.type == ptx.type.XTensorType("float64", dims=["x", "y"], shape=(3, 4))
+
+z = zx.values
+z.dprint()
+```
+
+
+```{testoutput}
+
+TensorFromXTensor [id A]
+ └─ XElemwise{scalar_op=Add()} [id B]
+    ├─ XTensorFromTensor{dims=('x',)} [id C]
+    │  └─ a [id D]
+    └─ XTensorFromTensor{dims=('y',)} [id E]
+       └─ b [id F]
+```
+
+Once we compile the graph, no XOps are left.
+
+```{testcode}
+
+import pytensor
+
+with pytensor.config.change_flags(optimizer_verbose=True):
+    fn = pytensor.function([a, b], z)
+
+```
+
+```{testoutput}
+
+rewriting: rewrite lower_elemwise replaces XElemwise{scalar_op=Add()}.0 of XElemwise{scalar_op=Add()}(XTensorFromTensor{dims=('x',)}.0, XTensorFromTensor{dims=('y',)}.0) with XTensorFromTensor{dims=('x', 'y')}.0 of XTensorFromTensor{dims=('x', 'y')}(Add.0)
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('x',)}.0) with a of None
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('y',)}.0) with b of None
+rewriting: rewrite useless_tensor_from_xtensor replaces TensorFromXTensor.0 of TensorFromXTensor(XTensorFromTensor{dims=('x', 'y')}.0) with Add.0 of Add(ExpandDims{axis=1}.0, ExpandDims{axis=0}.0)
+
+```
+
+```{testcode}
+
+fn.dprint()
+```
+
+```{testoutput}
+
+Add [id A] 2
+ ├─ ExpandDims{axis=1} [id B] 1
+ │  └─ a [id C]
+ └─ ExpandDims{axis=0} [id D] 0
+    └─ b [id E]
+```
+
+
+## Index
+
+:::{toctree}
+:maxdepth: 1
+
+module_functions
+math
+linalg
+random
+type
+:::

doc/library/xtensor/linalg.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_linalg)=
+# `xtensor.linalg` -- Linear algebra operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.linalg
+   :members:
+```

doc/library/xtensor/math.md

@@ -0,0 +1,8 @@
+(libdoc_xtensor_math)=
+# `xtensor.math` Mathematical operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.math
+   :members:
+   :exclude-members: XDot, dot
+```

doc/library/xtensor/module_functions.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_module_function)=
+# `xtensor` -- Module level operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor
+   :members: broadcast, concat, dot, full_like, ones_like, zeros_like
+```

doc/library/xtensor/random.md

@@ -0,0 +1,7 @@
+(libdoc_xtensor_random)=
+# `xtensor.random` Random number generator operations
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.random
+   :members:
+```

doc/library/xtensor/type.md

@@ -0,0 +1,20 @@
+(libdoc_xtenor_type)=
+
+# `xtensor.type` -- Types and Variables
+
+## XTensorVariable creation functions
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.type
+   :members: xtensor, xtensor_constant, as_xtensor
+
+```
+
+## XTensor Type and Variable classes
+
+```{eval-rst}
+.. automodule:: pytensor.xtensor.type
+   :members: XTensorType, XTensorVariable, XTensorConstant
+```
+
+