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

pytensor/link/jax/dispatch/slinalg.py

@@ -1,7 +1,30 @@
 import jax
 
 from pytensor.link.jax.dispatch.basic import jax_funcify
-from pytensor.tensor.slinalg import BlockDiagonal, Cholesky, Solve, SolveTriangular
+from pytensor.tensor.slinalg import (
+    BlockDiagonal,
+    Cholesky,
+    Eigvalsh,
+    Solve,
+    SolveTriangular,
+)
+
+
+@jax_funcify.register(Eigvalsh)
+def jax_funcify_Eigvalsh(op, **kwargs):
+    if op.lower:
+        UPLO = "L"
+    else:
+        UPLO = "U"
+
+    def eigvalsh(a, b):
+        if b is not None:
+            raise NotImplementedError(
+                "jax.numpy.linalg.eigvalsh does not support generalized eigenvector problems (b != None)"
+            )
+        return jax.numpy.linalg.eigvalsh(a, UPLO=UPLO)
+
+    return eigvalsh
 
 
 @jax_funcify.register(Cholesky)

pytensor/link/jax/dispatch/sort.py

@@ -0,0 +1,24 @@
+from jax import numpy as jnp
+
+from pytensor.link.jax.dispatch import jax_funcify
+from pytensor.tensor.sort import ArgSortOp, SortOp
+
+
+@jax_funcify.register(SortOp)
+def jax_funcify_Sort(op, **kwargs):
+    stable = op.kind == "stable"
+
+    def sort(arr, axis):
+        return jnp.sort(arr, axis=axis, stable=stable)
+
+    return sort
+
+
+@jax_funcify.register(ArgSortOp)
+def jax_funcify_ArgSort(op, **kwargs):
+    stable = op.kind == "stable"
+
+    def argsort(arr, axis):
+        return jnp.argsort(arr, axis=axis, stable=stable)
+
+    return argsort

pytensor/link/jax/dispatch/subtensor.py

@@ -31,36 +31,20 @@
 """
 
 
-def subtensor_assert_indices_jax_compatible(node, idx_list):
-    from pytensor.graph.basic import Constant
-    from pytensor.tensor.variable import TensorVariable
-
-    ilist = indices_from_subtensor(node.inputs[1:], idx_list)
-    for idx in ilist:
-        if isinstance(idx, TensorVariable):
-            if idx.type.dtype == "bool":
-                raise NotImplementedError(BOOLEAN_MASK_ERROR)
-        elif isinstance(idx, slice):
-            for slice_arg in (idx.start, idx.stop, idx.step):
-                if slice_arg is not None and not isinstance(slice_arg, Constant):
-                    raise NotImplementedError(DYNAMIC_SLICE_LENGTH_ERROR)
-
-
 @jax_funcify.register(Subtensor)
 @jax_funcify.register(AdvancedSubtensor)
 @jax_funcify.register(AdvancedSubtensor1)
 def jax_funcify_Subtensor(op, node, **kwargs):
     idx_list = getattr(op, "idx_list", None)
-    subtensor_assert_indices_jax_compatible(node, idx_list)
 
-    def subtensor_constant(x, *ilists):
+    def subtensor(x, *ilists):
         indices = indices_from_subtensor(ilists, idx_list)
         if len(indices) == 1:
             indices = indices[0]
 
         return x.__getitem__(indices)
 
-    return subtensor_constant
+    return subtensor
 
 
 @jax_funcify.register(IncSubtensor)

pytensor/link/jax/dispatch/tensor_basic.py

@@ -22,7 +22,6 @@
 )
 from pytensor.tensor.exceptions import NotScalarConstantError
 from pytensor.tensor.shape import Shape_i
-from pytensor.tensor.sort import SortOp
 
 
 ARANGE_CONCRETE_VALUE_ERROR = """JAX requires the arguments of `jax.numpy.arange`
@@ -206,11 +205,3 @@ def tri(*args):
         return jnp.tri(*args, dtype=op.dtype)
 
     return tri
-
-
-@jax_funcify.register(SortOp)
-def jax_funcify_Sort(op, **kwargs):
-    def sort(arr, axis):
-        return jnp.sort(arr, axis=axis)
-
-    return sort

pytensor/link/numba/dispatch/__init__.py

@@ -11,5 +11,6 @@
 import pytensor.link.numba.dispatch.scan
 import pytensor.link.numba.dispatch.sparse
 import pytensor.link.numba.dispatch.slinalg
+import pytensor.link.numba.dispatch.subtensor
 
 # isort: on