Add initial support for PyTorch backend

.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:
@@ -143,6 +149,7 @@ jobs:
           mamba 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 mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
           if [[ $INSTALL_JAX == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
+          if [[ $INSTALL_TORCH == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch && pip install tensorflow-probability; fi
           pip install -e ./
           mamba list && pip freeze
           python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
@@ -151,6 +158,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: bash -l {0}
@@ -195,7 +203,7 @@ jobs:
         - name: Install dependencies
           shell: bash -l {0}
           run: |
-            mamba 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
+            mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service cython pytest "numba>=0.57" jax jaxlib pytorch pytest-benchmark
             pip install -e ./
             mamba list && pip freeze
             python -c 'import pytensor; print(pytensor.config.__str__(print_doc=False))'
@@ -264,3 +272,4 @@ jobs:
           directory: ./coverage/
           fail_ci_if_error: true
           token: ${{ secrets.CODECOV_TOKEN }}
+

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(),
 }
 
@@ -462,6 +464,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(
@@ -476,6 +490,7 @@ def clone(self, link_kwargs=None, optimizer="", **kwargs):
     "FAST_RUN": FAST_RUN,
     "JAX": JAX,
     "NUMBA": NUMBA,
+    "PYTORCH": PYTORCH,
 }
 
 instantiated_default_mode = None

pytensor/link/pytorch/dispatch/__init__.py

@@ -0,0 +1,19 @@
+# isort: off
+from pytensor.link.pytorch.dispatch.basic import pytorch_funcify, pytorch_typify
+
+# # Load dispatch specializations
+import pytensor.link.pytorch.dispatch.scalar
+
+# import pytensor.link.jax.dispatch.tensor_basic
+# import pytensor.link.jax.dispatch.subtensor
+# import pytensor.link.jax.dispatch.shape
+# import pytensor.link.jax.dispatch.extra_ops
+# import pytensor.link.jax.dispatch.nlinalg
+# import pytensor.link.jax.dispatch.slinalg
+# import pytensor.link.jax.dispatch.random
+import pytensor.link.pytorch.dispatch.elemwise
+# import pytensor.link.jax.dispatch.scan
+# import pytensor.link.jax.dispatch.sparse
+# import pytensor.link.jax.dispatch.blockwise
+
+# isort: on

pytensor/link/pytorch/dispatch/basic.py

@@ -0,0 +1,105 @@
+import warnings
+from functools import singledispatch
+
+import torch
+
+from pytensor.compile.ops import DeepCopyOp, ViewOp
+from pytensor.graph.fg import FunctionGraph
+from pytensor.ifelse import IfElse
+from pytensor.link.utils import fgraph_to_python
+from pytensor.raise_op import CheckAndRaise
+
+
+@singledispatch
+def pytorch_typify(data, dtype=None, **kwargs):
+    r"""Convert instances of PyTensor `Type`\s to PyTorch types."""
+    if dtype is None:
+        return data
+    else:
+        return torch.tensor(data, dtype=dtype)
+
+
+@pytorch_typify.register(torch.Tensor)
+def pytorch_typify_tensor(data, dtype=None, **kwargs):
+    # if len(data.shape) == 0:
+    #     return data.item()
+    return torch.tensor(data, dtype=dtype)
+
+
+@singledispatch
+def pytorch_funcify(op, node=None, storage_map=None, **kwargs):
+    """Create a PyTorch compatible function from an PyTensor `Op`."""
+    raise NotImplementedError(f"No PyTorch conversion for the given `Op`: {op}")
+
+
+@pytorch_funcify.register(FunctionGraph)
+def pytorch_funcify_FunctionGraph(
+    fgraph,
+    node=None,
+    fgraph_name="pytorch_funcified_fgraph",
+    **kwargs,
+):
+    return fgraph_to_python(
+        fgraph,
+        pytorch_funcify,
+        type_conversion_fn=pytorch_typify,
+        fgraph_name=fgraph_name,
+        **kwargs,
+    )
+
+
+@pytorch_funcify.register(IfElse)
+def pytorch_funcify_IfElse(op, **kwargs):
+    n_outs = op.n_outs
+
+    def ifelse(cond, *args, n_outs=n_outs):
+        res = torch.where(
+            cond,
+            args[:n_outs][0],
+            args[n_outs:][0],
+        )
+        return res
+
+    return ifelse
+
+
+@pytorch_funcify.register(CheckAndRaise)
+def pytorch_funcify_CheckAndRaise(op, **kwargs):
+    def assert_fn(x, *conditions):
+        for cond in conditions:
+            assert cond.item()
+        return x
+
+    return assert_fn
+
+
+def pytorch_safe_copy(x):
+    try:
+        res = x.clone()
+    except NotImplementedError:
+        # warnings.warn(
+        #     "`jnp.copy` is not implemented yet. Using the object's `copy` method."
+        # )
+        if hasattr(x, "copy"):
+            res = torch.tensor(x.copy())
+        else:
+            warnings.warn(f"Object has no `copy` method: {x}")
+            res = x
+
+    return res
+
+
+@pytorch_funcify.register(DeepCopyOp)
+def pytorch_funcify_DeepCopyOp(op, **kwargs):
+    def deepcopyop(x):
+        return pytorch_safe_copy(x)
+
+    return deepcopyop
+
+
+@pytorch_funcify.register(ViewOp)
+def pytorch_funcify_ViewOp(op, **kwargs):
+    def viewop(x):
+        return x
+
+    return viewop

pytensor/link/pytorch/dispatch/elemwise.py

@@ -0,0 +1,68 @@
+import torch
+
+from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
+from pytensor.tensor.elemwise import DimShuffle, Elemwise
+from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
+
+
+@pytorch_funcify.register(Elemwise)
+def pytorch_funcify_Elemwise(op, node, **kwargs):
+    scalar_op = op.scalar_op
+    base_fn = pytorch_funcify(scalar_op, node=node, **kwargs)
+
+    def elemwise_fn(*inputs):
+        # Elemwise._check_runtime_broadcast(node, tuple(map(torch.tensor, inputs)))
+        return base_fn(*inputs)
+
+    return elemwise_fn
+
+
+@pytorch_funcify.register(DimShuffle)
+def pytorch_funcify_DimShuffle(op, **kwargs):
+    def dimshuffle(x):
+        res = torch.transpose(x, *op.transposition)
+
+        shape = list(res.shape[: len(op.shuffle)])
+
+        for augm in op.augment:
+            shape.insert(augm, 1)
+
+        res = torch.reshape(res, shape)
+
+        if not op.inplace:
+            res = res.clone()
+
+        return res
+
+    return dimshuffle
+
+
+@pytorch_funcify.register(Softmax)
+def pytorch_funcify_Softmax(op, **kwargs):
+    axis = op.axis
+
+    def softmax(x):
+        return torch.nn.functional.softmax(x, dim=axis)
+
+    return softmax
+
+
+@pytorch_funcify.register(SoftmaxGrad)
+def pytorch_funcify_SoftmaxGrad(op, **kwargs):
+    axis = op.axis
+
+    def softmax_grad(dy, sm):
+        dy_times_sm = dy * sm
+        return dy_times_sm - torch.sum(dy_times_sm, dim=axis, keepdims=True) * sm
+
+    return softmax_grad
+
+
+@pytorch_funcify.register(LogSoftmax)
+def pytorch_funcify_LogSoftmax(op, **kwargs):
+    axis = op.axis
+
+    def log_softmax(x):
+        return torch.nn.functional.log_softmax(x, dim=axis)
+
+    return log_softmax