Merge branch 'pymc-devs:main' into main

Author: Prtm2110

.github/workflows/pypi.yml

@@ -49,7 +49,7 @@ jobs:
           fetch-depth: 0
 
       - name: Build wheels
-        uses: pypa/cibuildwheel@v2.18.1
+        uses: pypa/cibuildwheel@v2.19.1
 
       - uses: actions/upload-artifact@v3
         with:
@@ -88,7 +88,7 @@ jobs:
         name: artifact
         path: dist
 
-    - uses: pypa/gh-action-pypi-publish@v1.8.14
+    - uses: pypa/gh-action-pypi-publish@v1.9.0
       with:
         user: __token__
         password: ${{ secrets.pypi_password }}

.github/workflows/test.yml

@@ -76,6 +76,7 @@ jobs:
         float32: [0, 1]
         install-numba: [0]
         install-jax: [0]
+        install-torch: [0]
         part:
           - "tests --ignore=tests/tensor --ignore=tests/scan --ignore=tests/sparse"
           - "tests/scan"
@@ -116,6 +117,11 @@ jobs:
             fast-compile: 0
             float32: 0
             part: "tests/link/jax"
+          - install-torch: 1
+            python-version: "3.10"
+            fast-compile: 0
+            float32: 0
+            part: "tests/link/pytorch"
     steps:
       - uses: actions/checkout@v4
         with:
@@ -142,9 +148,12 @@ jobs:
       - name: Install dependencies
         shell: micromamba-shell {0}
         run: |
+
           micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock sympy
           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
+
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
@@ -153,6 +162,7 @@ jobs:
           PYTHON_VERSION: ${{ matrix.python-version }}
           INSTALL_NUMBA: ${{ matrix.install-numba }}
           INSTALL_JAX: ${{ matrix.install-jax }}
+          INSTALL_TORCH: ${{ matrix.install-torch}}
 
       - name: Run tests
         shell: micromamba-shell {0}
@@ -199,7 +209,7 @@ jobs:
       - 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
+          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
           pip install -e ./
           micromamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
@@ -268,3 +278,4 @@ jobs:
           directory: ./coverage/
           fail_ci_if_error: true
           token: ${{ secrets.CODECOV_TOKEN }}
+

.pre-commit-config.yaml

@@ -22,7 +22,7 @@ repos:
           )$
       - id: check-merge-conflict
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.4.8
+    rev: v0.4.10
     hooks:
       - id: ruff
         args: ["--fix", "--output-format=full"]

environment.yml

@@ -9,8 +9,8 @@ channels:
 dependencies:
   - python>=3.10
   - compilers
-  - numpy>=1.17.0
-  - scipy>=0.14
+  - numpy>=1.17.0,<2
+  - scipy>=0.14,<1.14.0
   - filelock
   - etuples
   - logical-unification

pyproject.toml

@@ -47,7 +47,7 @@ keywords = [
 ]
 dependencies = [
     "setuptools>=59.0.0",
-    "scipy>=0.14",
+    "scipy>=0.14,<1.14",
     "numpy>=1.17.0,<2",
     "filelock",
     "etuples",

pytensor/__init__.py

@@ -108,7 +108,7 @@ def as_symbolic(x: Any, name: str | None = None, **kwargs) -> Variable:
 
 
 @singledispatch
-def _as_symbolic(x, **kwargs) -> Variable:
+def _as_symbolic(x: Any, **kwargs) -> Variable:
     from pytensor.tensor import as_tensor_variable
 
     return as_tensor_variable(x, **kwargs)

pytensor/compile/mode.py

@@ -28,6 +28,7 @@
 from pytensor.link.c.basic import CLinker, OpWiseCLinker
 from pytensor.link.jax.linker import JAXLinker
 from pytensor.link.numba.linker import NumbaLinker
+from pytensor.link.pytorch.linker import PytorchLinker
 from pytensor.link.vm import VMLinker
 
 
@@ -47,6 +48,7 @@
     "vm_nogc": VMLinker(allow_gc=False, use_cloop=False),
     "cvm_nogc": VMLinker(allow_gc=False, use_cloop=True),
     "jax": JAXLinker(),
+    "pytorch": PytorchLinker(),
     "numba": NumbaLinker(),
 }
 
@@ -460,6 +462,18 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
         ],
     ),
 )
+PYTORCH = Mode(
+    PytorchLinker(),
+    RewriteDatabaseQuery(
+        include=["fast_run"],
+        exclude=[
+            "cxx_only",
+            "BlasOpt",
+            "fusion",
+            "inplace",
+        ],
+    ),
+)
 NUMBA = Mode(
     NumbaLinker(),
     RewriteDatabaseQuery(
@@ -474,6 +488,7 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     "FAST_RUN": FAST_RUN,
     "JAX": JAX,
     "NUMBA": NUMBA,
+    "PYTORCH": PYTORCH,
 }
 
 instantiated_default_mode = None

pytensor/graph/basic.py

@@ -1302,8 +1302,8 @@ def clone_node_and_cache(
 
 
 def clone_get_equiv(
-    inputs: Sequence[Variable],
-    outputs: Sequence[Variable],
+    inputs: Iterable[Variable],
+    outputs: Reversible[Variable],
     copy_inputs: bool = True,
     copy_orphans: bool = True,
     memo: dict[Union[Apply, Variable, "Op"], Union[Apply, Variable, "Op"]]

pytensor/link/basic.py

@@ -600,6 +600,10 @@ def create_thunk_inputs(self, storage_map: dict[Variable, list[Any]]) -> list[An
     def jit_compile(self, fn: Callable) -> Callable:
         """JIT compile a converted ``FunctionGraph``."""
 
+    def input_filter(self, inp: Any) -> Any:
+        """Apply a filter to the data input."""
+        return inp
+
     def output_filter(self, var: Variable, out: Any) -> Any:
         """Apply a filter to the data output by a JITed function call."""
         return out
@@ -657,7 +661,7 @@ def thunk(
             thunk_inputs=thunk_inputs,
             thunk_outputs=thunk_outputs,
         ):
-            outputs = fgraph_jit(*[x[0] for x in thunk_inputs])
+            outputs = fgraph_jit(*[self.input_filter(x[0]) for x in thunk_inputs])
 
             for o_var, o_storage, o_val in zip(fgraph.outputs, thunk_outputs, outputs):
                 compute_map[o_var][0] = True

pytensor/link/jax/dispatch/__init__.py

@@ -14,5 +14,6 @@
 import pytensor.link.jax.dispatch.scan
 import pytensor.link.jax.dispatch.sparse
 import pytensor.link.jax.dispatch.blockwise
+import pytensor.link.jax.dispatch.sort
 
 # isort: on