Add docs on implementing Pytorch Ops (and CumOp)

pytensor/link/pytorch/dispatch/__init__.py

@@ -4,4 +4,5 @@
 # # Load dispatch specializations
 import pytensor.link.pytorch.dispatch.scalar
 import pytensor.link.pytorch.dispatch.elemwise
+import pytensor.link.pytorch.dispatch.extra_ops
 # isort: on

pytensor/link/pytorch/dispatch/extra_ops.py

@@ -0,0 +1,21 @@
+import torch
+
+from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
+from pytensor.tensor.extra_ops import CumOp
+
+
+@pytorch_funcify.register(CumOp)
+def pytorch_funcify_Cumop(op, **kwargs):
+    dim = op.axis
+    mode = op.mode
+
+    def cumop(x, dim=dim):
+        if dim is None:
+            x = x.reshape(-1)
+            dim = 0
+        if mode == "add":
+            return torch.cumsum(x, dim=dim)
+        else:
+            return torch.cumprod(x, dim=dim)
+
+    return cumop

tests/link/pytorch/test_extra_ops.py

@@ -0,0 +1,38 @@
+import numpy as np
+import pytest
+
+import pytensor.tensor as pt
+from pytensor.configdefaults import config
+from pytensor.graph import FunctionGraph
+from tests.link.pytorch.test_basic import compare_pytorch_and_py
+
+
+@pytest.mark.parametrize(
+    "axis",
+    [
+        None,
+        1,
+    ],
+)
+def test_pytorch_CumOp(axis):
+    """Test PyTorch conversion of the `CumOp` `Op`."""
+
+    # Create a symbolic input for the first input of `CumOp`
+    a = pt.matrix("a")
+
+    # Create test value tag for a
+    test_value = np.arange(9, dtype=config.floatX).reshape((3, 3))
+
+    # Create the output variable
+    out = pt.cumsum(a, axis=axis)
+
+    # Create a PyTensor `FunctionGraph`
+    fgraph = FunctionGraph([a], [out])
+
+    # Pass the graph and inputs to the testing function
+    compare_pytorch_and_py(fgraph, [test_value])
+
+    # For the second mode of CumOp
+    out = pt.cumprod(a, axis=axis)
+    fgraph = FunctionGraph([a], [out])
+    compare_pytorch_and_py(fgraph, [test_value])