Pytorch support for Join and Careduce Ops

pytensor/link/pytorch/dispatch/elemwise.py

@@ -41,7 +41,7 @@
 @pytorch_funcify.register(Sum)
 def pytorch_funcify_sum(op, **kwargs):
     def torch_sum(x):
         return torch.sum(x, dim=op.axis)
 
     return torch_sum
 
@@ -49,7 +49,7 @@
 @pytorch_funcify.register(All)
 def pytorch_funcify_all(op, **kwargs):
     def torch_all(x):
         return torch.all(x, dim=op.axis)
 
     return torch_all
 
@@ -58,11 +58,11 @@
 def pytorch_funcify_prod(op, **kwargs):
     def torch_prod(x):
         if isinstance(op.axis, tuple):
-            for d in op.axis:
-                x = torch.prod(x, dim=d, keepdim=True)
-            return x.squeeze()
+            for d in op.axis[::-1]:
+                x = torch.prod(x, dim=d)
+            return x
         else:
             return torch.prod(x.flatten(), dim=0)
 
     return torch_prod
 
@@ -70,7 +70,7 @@
 @pytorch_funcify.register(Any)
 def pytorch_funcify_any(op, **kwargs):
     def torch_any(x):
         return torch.any(x, dim=op.axis)
 
     return torch_any
 
@@ -79,11 +79,11 @@
 def pytorch_funcify_max(op, **kwargs):
     def torch_max(x):
         if isinstance(op.axis, tuple):
-            for d in op.axis:
-                x = torch.max(x, dim=d, keepdim=True).values
-            return x.squeeze()
+            for d in op.axis[::-1]:
+                x = torch.max(x, dim=d).values
+            return x
         else:
             return torch.max(x.flatten(), dim=0).values
 
     return torch_max
 
@@ -92,11 +92,11 @@
 def pytorch_funcify_min(op, **kwargs):
     def torch_min(x):
         if isinstance(op.axis, tuple):
-            for d in op.axis:
-                x = torch.min(x, dim=d, keepdim=True).values
-            return x.squeeze()
+            for d in op.axis[::-1]:
+                x = torch.min(x, dim=d).values
+            return x
         else:
             return torch.min(x.flatten(), dim=0).values
 
     return torch_min
 

tests/link/pytorch/test_elemwise.py

@@ -58,8 +58,9 @@ def test_pytorch_elemwise():
     compare_pytorch_and_py(fg, [[0.9, 0.9]])
 
 
+@pytest.mark.parametrize("fn", [ptm.sum, ptm.prod, ptm.max, ptm.min])
 @pytest.mark.parametrize("axis", [0, 1, (0, 1), (1, 2), (1, -1)])
-def test_pytorch_sum(axis):
+def test_pytorch_careduce(fn, axis):
     a_pt = tensor3("a")
     test_value = np.array(
         [
@@ -79,113 +80,19 @@ def test_pytorch_sum(axis):
         ]
     ).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", [None, 0, 1, (0, 1), (1, 2)])
-def test_pytorch_prod(axis):
-    a_pt = tensor3("a")
-    test_value = np.array(
-        [
-            [
-                [1, 1, 1, 1],
-                [2, 2, 2, 2],
-            ],
-            [
-                [3, 3, 3, 3],
-                [
-                    4,
-                    4,
-                    4,
-                    4,
-                ],
-            ],
-        ]
-    ).astype(config.floatX)
-
-    x = ptm.prod(a_pt, axis=axis)
+    x = fn(a_pt, axis=axis)
     x_fg = FunctionGraph([a_pt], [x])
 
     compare_pytorch_and_py(x_fg, [test_value])
 
 
+@pytest.mark.parametrize("fn", [ptm.any, ptm.all])
 @pytest.mark.parametrize("axis", [None, 0, 1, (0, 1)])
-def test_pytorch_any(axis):
+def test_pytorch_any_all(fn, 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", [None, 0, 1, (1, -1)])
-def test_pytorch_max(axis):
-    a_pt = tensor3("a")
-    test_value = np.array(
-        [
-            [
-                [1, 1, 1, 1],
-                [2, 2, 2, 2],
-            ],
-            [
-                [3, 3, 3, 3],
-                [
-                    4,
-                    4,
-                    4,
-                    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", [None, 0, 1, (1, -1)])
-def test_pytorch_min(axis):
-    a_pt = tensor3("a")
-    test_value = np.array(
-        [
-            [[1, 1, 1, 1], [2, 2, 2, 2], [3, 3, 3, 3]],
-            [
-                [4, 4, 4, 4],
-                [
-                    5,
-                    5,
-                    5,
-                    5,
-                ],
-                [
-                    6,
-                    6,
-                    6,
-                    6,
-                ],
-            ],
-        ]
-    ).astype(config.floatX)
-
-    x = ptm.min(a_pt, axis=axis)
+    x = fn(a_pt, axis=axis)
     x_fg = FunctionGraph([a_pt], [x])
 
     compare_pytorch_and_py(x_fg, [test_value])