Pytorch support for Join and Careduce Ops

pytensor/link/pytorch/dispatch/basic.py

@@ -6,7 +6,7 @@
 from pytensor.graph.fg import FunctionGraph
 from pytensor.link.utils import fgraph_to_python
 from pytensor.raise_op import CheckAndRaise
-from pytensor.tensor.basic import Alloc, AllocEmpty, ARange
+from pytensor.tensor.basic import Alloc, AllocEmpty, ARange, Join
 
 
 @singledispatch
@@ -89,3 +89,14 @@ def arange(start, stop, step):
         return torch.arange(start, stop, step, dtype=dtype)
 
     return arange
+
+
+@pytorch_funcify.register(Join)
+def pytorch_funcify_Join(op, **kwargs):
+    def join(axis, *tensors):
+        # tensors could also be tuples, and in this case they don't have a ndim
+        tensors = [torch.tensor(tensor) for tensor in tensors]
+
+        return torch.cat(tensors, dim=axis)
+
+    return join

pytensor/link/pytorch/dispatch/elemwise.py

@@ -2,6 +2,7 @@
 
 from pytensor.link.pytorch.dispatch.basic import pytorch_funcify
 from pytensor.tensor.elemwise import DimShuffle, Elemwise
+from pytensor.tensor.math import All, Any, Max, Min, Prod, Sum
 from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
 
 
@@ -37,6 +38,64 @@
     return dimshuffle
 
 
+@pytorch_funcify.register(Sum)
+def pytorch_funcify_sum(op, **kwargs):
+    def torch_sum(x):
+        return torch.sum(x, dim=op.axis)
+
+    return torch_sum
+
+
+@pytorch_funcify.register(All)
+def pytorch_funcify_all(op, **kwargs):
+    dim = op.axis
+
+    def torch_all(x):
+        return torch.all(x, dim=dim)
+
+    return torch_all
+
+
+@pytorch_funcify.register(Prod)
+def pytorch_funcify_prod(op, **kwargs):
+    dim = op.axis[0]
+
+    def torch_prod(x):
+        return torch.prod(x, dim=dim)
+
+    return torch_prod
+
+
+@pytorch_funcify.register(Any)
+def pytorch_funcify_any(op, **kwargs):
+    dim = op.axis
+
+    def torch_any(x):
+        return torch.any(x, dim=dim)
+
+    return torch_any
+
+
+@pytorch_funcify.register(Max)
+def pytorch_funcify_max(op, **kwargs):
+    dim = op.axis[0]
+
+    def torch_max(x):
+        return torch.max(x, dim=dim).values
+
+    return torch_max
+
+
+@pytorch_funcify.register(Min)
+def pytorch_funcify_min(op, **kwargs):
+    dim = op.axis[0]
+
+    def torch_min(x):
+        return torch.min(x, dim=dim).values
+
+    return torch_min
+
+
 @pytorch_funcify.register(Softmax)
 def pytorch_funcify_Softmax(op, **kwargs):
     axis = op.axis

tests/link/pytorch/test_basic.py

@@ -4,6 +4,7 @@
 import numpy as np
 import pytest
 
+import pytensor.tensor.basic as ptb
 from pytensor.compile.function import function
 from pytensor.compile.mode import get_mode
 from pytensor.compile.sharedvalue import SharedVariable, shared
@@ -13,7 +14,7 @@
 from pytensor.graph.op import Op
 from pytensor.raise_op import CheckAndRaise
 from pytensor.tensor import alloc, arange, as_tensor, empty
-from pytensor.tensor.type import scalar, vector
+from pytensor.tensor.type import matrix, scalar, vector
 
 
 torch = pytest.importorskip("torch")
@@ -235,3 +236,42 @@ def test_arange():
         FunctionGraph([start, stop, step], [out]),
         [np.array(1), np.array(10), np.array(2)],
     )
+
+
+def test_pytorch_Join():
+    a = matrix("a")
+    b = matrix("b")
+
+    x = ptb.join(0, a, b)
+    x_fg = FunctionGraph([a, b], [x])
+    compare_pytorch_and_py(
+        x_fg,
+        [
+            np.c_[[1.0, 2.0, 3.0]].astype(config.floatX),
+            np.c_[[4.0, 5.0, 6.0]].astype(config.floatX),
+        ],
+    )
+    compare_pytorch_and_py(
+        x_fg,
+        [
+            np.c_[[1.0, 2.0, 3.0]].astype(config.floatX),
+            np.c_[[4.0, 5.0]].astype(config.floatX),
+        ],
+    )
+
+    x = ptb.join(1, a, b)
+    x_fg = FunctionGraph([a, b], [x])
+    compare_pytorch_and_py(
+        x_fg,
+        [
+            np.c_[[1.0, 2.0, 3.0]].astype(config.floatX),
+            np.c_[[4.0, 5.0, 6.0]].astype(config.floatX),
+        ],
+    )
+    compare_pytorch_and_py(
+        x_fg,
+        [
+            np.c_[[1.0, 2.0], [3.0, 4.0]].astype(config.floatX),
+            np.c_[[5.0, 6.0]].astype(config.floatX),
+        ],
+    )

tests/link/pytorch/test_elemwise.py

@@ -2,6 +2,7 @@
 import pytest
 
 import pytensor.tensor as pt
+import pytensor.tensor.math as ptm
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
 from pytensor.tensor import elemwise as pt_elemwise
@@ -57,6 +58,72 @@ def test_pytorch_elemwise():
     compare_pytorch_and_py(fg, [[0.9, 0.9]])
 
 
+@pytest.mark.parametrize("axis", [None, 0, 1])
+def test_pytorch_sum(axis):
+    a_pt = matrix("a")
+    test_value = np.array([[1, 2], [3, 4]]).astype(config.floatX)
+
+    x = pt.math.sum(a_pt, axis=axis)
+    x_fg = FunctionGraph([a_pt], [x])
+
+    compare_pytorch_and_py(x_fg, [test_value])
+
+
+@pytest.mark.parametrize("axis", [None, 0, 1])
+def test_pytorch_all(axis):
+    a_pt = matrix("a")
+    test_value = np.array([[True, False, True], [False, True, True]])
+
+    x = ptm.all(a_pt, axis=axis)
+    x_fg = FunctionGraph([a_pt], [x])
+
+    compare_pytorch_and_py(x_fg, [test_value])
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+def test_pytorch_prod(axis):
+    a_pt = matrix("a")
+    test_value = np.array([[1, 2], [3, 4]]).astype(config.floatX)
+
+    x = ptm.prod(a_pt, axis=axis)
+    x_fg = FunctionGraph([a_pt], [x])
+
+    compare_pytorch_and_py(x_fg, [test_value])
+
+
+@pytest.mark.parametrize("axis", [None, 0, 1])
+def test_pytorch_any(axis):
+    a_pt = matrix("a")
+    test_value = np.array([[True, False, True], [False, True, True]])
+
+    x = ptm.any(a_pt, axis=axis)
+    x_fg = FunctionGraph([a_pt], [x])
+
+    compare_pytorch_and_py(x_fg, [test_value])
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+def test_pytorch_max(axis):
+    a_pt = matrix("a")
+    test_value = np.array([[1, 2], [3, 4]]).astype(config.floatX)
+
+    x = ptm.max(a_pt, axis=axis)
+    x_fg = FunctionGraph([a_pt], [x])
+
+    compare_pytorch_and_py(x_fg, [test_value])
+
+
+@pytest.mark.parametrize("axis", [0, 1])
+def test_pytorch_min(axis):
+    a_pt = matrix("a")
+    test_value = np.array([[1, 2], [3, 4]]).astype(config.floatX)
+
+    x = ptm.min(a_pt, axis=axis)
+    x_fg = FunctionGraph([a_pt], [x])
+
+    compare_pytorch_and_py(x_fg, [test_value])
+
+
 @pytest.mark.parametrize("dtype", ["float64", "int64"])
 @pytest.mark.parametrize("axis", [None, 0, 1])
 def test_softmax(axis, dtype):