Merge branch 'main' of github.com:pymc-devs/pytensor into blockwise

Author: Ch0ronomato

.github/workflows/mypy.yml

@@ -15,7 +15,7 @@ jobs:
       - uses: actions/checkout@v4
       - uses: mamba-org/setup-micromamba@v1
         with:
-          micromamba-version: "latest" # any version from https://github.com/mamba-org/micromamba-releases
+          micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
           environment-file: environment.yml
           init-shell: bash
           cache-environment: true

.github/workflows/pypi.yml

@@ -35,6 +35,18 @@ jobs:
           name: sdist
           path: dist/*.tar.gz
 
+  run_checks:
+    name: Build & inspect our package.
+    # Note: the resulting builds are not actually published.
+    # This is purely for additional testing and diagnostic purposes.
+    runs-on: ubuntu-latest
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+      - uses: hynek/build-and-inspect-python-package@v2
+
   build_wheels:
     name: Build wheels for ${{ matrix.platform }}
     runs-on: ${{ matrix.platform }}
@@ -50,7 +62,7 @@ jobs:
           fetch-depth: 0
 
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.19.2
+        uses: pypa/cibuildwheel@v2.21.2
 
       - uses: actions/upload-artifact@v4
         with:
@@ -133,7 +145,7 @@ jobs:
           name: universal_wheel
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@v1.9.0
+      - uses: pypa/gh-action-pypi-publish@v1.10.3
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

.github/workflows/rtd-link-preview.yml

@@ -0,0 +1,16 @@
+name: Read the Docs Pull Request Preview
+on:
+  pull_request_target:
+    types:
+      - opened
+
+permissions:
+  pull-requests: write
+
+jobs:
+  documentation-links:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: readthedocs/actions/preview@v1
+        with:
+          project-slug: "pytensor"

.github/workflows/test.yml

@@ -78,7 +78,7 @@ jobs:
         install-jax: [0]
         install-torch: [0]
         part:
-          - "--doctest-modules --ignore=pytensor/misc/check_duplicate_key.py pytensor --ignore=pytensor/link"
+          - "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
           - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse"
           - "tests/scan"
           - "tests/sparse"
@@ -97,9 +97,9 @@ jobs:
             part: "tests/tensor/test_math.py"
           - fast-compile: 1
             float32: 1
-          - part: "--doctest-modules --ignore=pytensor/misc/check_duplicate_key.py pytensor --ignore=pytensor/link"
+          - part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
             float32: 1
-          - part: "--doctest-modules --ignore=pytensor/misc/check_duplicate_key.py pytensor --ignore=pytensor/link"
+          - part: "--doctest-modules pytensor --ignore=pytensor/misc/check_duplicate_key.py --ignore=pytensor/link"
             fast-compile: 1
         include:
           - install-numba: 1
@@ -135,7 +135,7 @@ jobs:
         uses: mamba-org/setup-micromamba@v1
         with:
           environment-name: pytensor-test
-          micromamba-version: "latest"
+          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 }}
@@ -157,7 +157,7 @@ jobs:
           micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
           if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
-          if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 -c pytorch -c nvidia; fi
+          if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 "mkl<=2024.0" -c pytorch -c nvidia; fi
           pip install pytest-sphinx
           
           pip install -e ./
@@ -209,13 +209,13 @@ jobs:
         uses: mamba-org/setup-micromamba@v1
         with:
           environment-name: pytensor-test
-          micromamba-version: "latest"
+          micromamba-version: "1.5.10-0" # until https://github.com/mamba-org/setup-micromamba/issues/225 is resolved
           init-shell: bash
           post-cleanup: "all"
       - name: Install dependencies
         shell: micromamba-shell {0}
         run: |
-          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark  pytorch pytorch-cuda=12.1 -c pytorch -c nvidia
+          micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytest-benchmark
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'

.pre-commit-config.yaml

@@ -9,7 +9,7 @@ exclude: |
     )$
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v5.0.0
     hooks:
       - id: debug-statements
         exclude: |
@@ -21,8 +21,13 @@ repos:
               pytensor/tensor/variable\.py|
           )$
       - id: check-merge-conflict
+  - repo: https://github.com/sphinx-contrib/sphinx-lint                                  
+    rev: v1.0.0                                                                          
+    hooks:                                                                               
+    - id: sphinx-lint 
+      args: ["."]
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.6.3
+    rev: v0.7.1
     hooks:
       - id: ruff
         args: ["--fix", "--output-format=full"]

readthedocs.yml renamed to .readthedocs.yaml

@@ -4,6 +4,6 @@ sphinx:
 conda:
   environment: doc/environment.yml
 build:
-  os: "ubuntu-20.04"
+  os: "ubuntu-lts-latest"
   tools:
-    python: "mambaforge-4.10"
+    python: "mambaforge-latest"

doc/extending/creating_a_c_op.rst

@@ -152,7 +152,7 @@ This distance between consecutive elements of an array over a given dimension,
 is called the stride of that dimension.
 
 
-Accessing NumPy :class`ndarray`\s' data and properties
+Accessing NumPy :class:`ndarray`'s data and properties
 ------------------------------------------------------
 
 The following macros serve to access various attributes of NumPy :class:`ndarray`\s.

doc/extending/creating_a_numba_jax_op.rst

@@ -4,7 +4,7 @@ Adding JAX, Numba and Pytorch support for `Op`\s
 PyTensor is able to convert its graphs into JAX, Numba and Pytorch compiled functions. In order to do
 this, each :class:`Op` in an PyTensor graph must have an equivalent JAX/Numba/Pytorch implementation function.
 
-This tutorial will explain how JAX, Numba and Pytorch implementations are created for an :class:`Op`. 
+This tutorial will explain how JAX, Numba and Pytorch implementations are created for an :class:`Op`.
 
 Step 1: Identify the PyTensor :class:`Op` you'd like to implement
 ------------------------------------------------------------------------
@@ -60,7 +60,7 @@ could also have any data type (e.g. floats, ints), so our implementation
 must be able to handle all the possible data types.
 
 It also tells us that there's only one return value, that it has a data type
-determined by :meth:`x.type()` i.e., the data type of the original tensor.
+determined by :meth:`x.type` i.e., the data type of the original tensor.
 This implies that the result is necessarily a matrix.
 
 Some class may have a more complex behavior. For example, the :class:`CumOp`\ :class:`Op`
@@ -116,7 +116,7 @@ Here's an example for :class:`DimShuffle`:
 
 .. tab-set::
 
-        .. tab-item:: JAX     
+        .. tab-item:: JAX
 
             .. code:: python
 
@@ -134,7 +134,7 @@ Here's an example for :class:`DimShuffle`:
                         res = jnp.copy(res)
 
                     return res
-        
+
         .. tab-item:: Numba
 
             .. code:: python
@@ -465,7 +465,7 @@ Step 4: Write tests
     .. tab-item:: JAX
 
         Test that your registered `Op` is working correctly by adding tests to the
-        appropriate test suites in PyTensor (e.g. in ``tests.link.jax``). 
+        appropriate test suites in PyTensor (e.g. in ``tests.link.jax``).
         The tests should ensure that your implementation can
         handle the appropriate types of inputs and produce outputs equivalent to `Op.perform`.
         Check the existing tests for the general outline of these kinds of tests. In
@@ -478,7 +478,7 @@ Step 4: Write tests
         Here's a small example of a test for :class:`CumOp` above:
 
         .. code:: python
-            
+
             import numpy as np
             import pytensor.tensor as pt
             from pytensor.configdefaults import config
@@ -514,22 +514,22 @@ Step 4: Write tests
         .. code:: python
 
             import pytest
-            
+
             def test_jax_CumOp():
                 """Test JAX conversion of the `CumOp` `Op`."""
                 a = pt.matrix("a")
                 a.tag.test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
-                
+
                 with pytest.raises(NotImplementedError):
                     out = pt.cumprod(a, axis=1)
                     fgraph = FunctionGraph([a], [out])
                     compare_jax_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
-    
-    
+
+
     .. tab-item:: Numba
 
         Test that your registered `Op` is working correctly by adding tests to the
-        appropriate test suites in PyTensor (e.g. in ``tests.link.numba``). 
+        appropriate test suites in PyTensor (e.g. in ``tests.link.numba``).
         The tests should ensure that your implementation can
         handle the appropriate types of inputs and produce outputs equivalent to `Op.perform`.
         Check the existing tests for the general outline of these kinds of tests. In
@@ -542,7 +542,7 @@ Step 4: Write tests
         Here's a small example of a test for :class:`CumOp` above:
 
         .. code:: python
-            
+
             from tests.link.numba.test_basic import compare_numba_and_py
             from pytensor.graph import FunctionGraph
             from pytensor.compile.sharedvalue import SharedVariable
@@ -561,11 +561,11 @@ Step 4: Write tests
                         if not isinstance(i, SharedVariable | Constant)
                     ],
                 )
-    
+
 
 
     .. tab-item:: Pytorch
-        
+
         Test that your registered `Op` is working correctly by adding tests to the
         appropriate test suites in PyTensor (``tests.link.pytorch``). The tests should ensure that your implementation can
         handle the appropriate types of inputs and produce outputs equivalent to `Op.perform`.
@@ -579,7 +579,7 @@ Step 4: Write tests
         Here's a small example of a test for :class:`CumOp` above:
 
         .. code:: python
-            
+
             import numpy as np
             import pytest
             import pytensor.tensor as pt
@@ -592,7 +592,7 @@ Step 4: Write tests
                 ["float64", "int64"],
             )
             @pytest.mark.parametrize(
-                "axis", 
+                "axis",
                 [None, 1, (0,)],
             )
             def test_pytorch_CumOp(axis, dtype):
@@ -650,4 +650,4 @@ as reported in issue `#654 <https://github.com/pymc-devs/pytensor/issues/654>`_.
 All jitted functions now must have constant shape, which means a graph like the
 one of :class:`Eye` can never be translated to JAX, since it's fundamentally a
 function with dynamic shapes. In other words, only PyTensor graphs with static shapes
-can be translated to JAX at the moment.
+can be translated to JAX at the moment.

doc/extending/type.rst

@@ -333,7 +333,7 @@ returns eitehr a new transferred variable (which can be the same as
 the input if no transfer is necessary) or returns None if the transfer
 can't be done.
 
-Then register that function by calling :func:`register_transfer()`
+Then register that function by calling :func:`register_transfer`
 with it as argument.
 
 An example

doc/library/compile/io.rst

@@ -36,7 +36,7 @@ The ``inputs`` argument to ``pytensor.function`` is a list, containing the ``Var
       ``self.<name>``. The default value is ``None``.
 
       ``value``: literal or ``Container``. The initial/default value for this
-        input. If update is`` None``, this input acts just like
+        input. If update is ``None``, this input acts just like
         an argument with a default value in Python. If update is not ``None``,
         changes to this
         value will "stick around", whether due to an update or a user's

doc/library/config.rst

@@ -226,7 +226,7 @@ import ``pytensor`` and print the config variable, as in:
     in the future.
 
     The ``'numpy+floatX'`` setting attempts to mimic NumPy casting rules,
-    although it prefers to use ``float32` `numbers instead of ``float64`` when
+    although it prefers to use ``float32`` numbers instead of ``float64`` when
     ``config.floatX`` is set to ``'float32'`` and the associated data is not
     explicitly typed as ``float64`` (e.g. regular Python floats).  Note that
     ``'numpy+floatX'`` is not currently behaving exactly as planned (it is a

doc/library/graph/graph.rst

@@ -4,12 +4,5 @@
 :mod:`graph` -- Interface for the PyTensor graph
 ================================================
 
----------
-Reference
----------
-
 .. automodule:: pytensor.graph.basic
-   :platform: Unix, Windows
-   :synopsis: Interface for types of symbolic variables
    :members:
-.. moduleauthor:: LISA

doc/library/graph/index.rst

@@ -1,20 +1,20 @@
 
 .. _libdoc_graph:
 
-================================================
-:mod:`graph` -- Theano Internals [doc TODO]
-================================================
+========================================
+:mod:`graph` -- PyTensor Graph Internals
+========================================
 
 .. module:: graph
-   :platform: Unix, Windows
-   :synopsis: Theano Internals
+
 .. moduleauthor:: LISA
 
 .. toctree::
     :maxdepth: 1
 
     graph
     fgraph
+    replace
     features
     op
     type

doc/library/graph/op.rst

@@ -1,12 +1,8 @@
-
 .. _libdoc_graph_op:
 
-==============================================================
-:mod:`graph` -- Objects and functions for computational graphs
-==============================================================
+===========================================
+:mod:`op` -- Objects that define operations
+===========================================
 
 .. automodule:: pytensor.graph.op
-   :platform: Unix, Windows
-   :synopsis: Interface for types of symbolic variables
    :members:
-.. moduleauthor:: LISA

doc/library/graph/replace.rst

@@ -0,0 +1,8 @@
+.. _libdoc_graph_replace:
+
+==================================================
+:mod:`replace` -- High level graph transformations
+==================================================
+
+.. automodule:: pytensor.graph.replace
+   :members: