Added PyTorch links for Max and Argmax. Tests are failing

Author: HangenYuu

pytensor/link/pytorch/dispatch/math.py

@@ -1,7 +1,8 @@
+import numpy as np
 import torch
 
 from pytensor.link.pytorch.dispatch import pytorch_funcify
-from pytensor.tensor.math import Dot
+from pytensor.tensor.math import Argmax, Dot, Max
 
 
 @pytorch_funcify.register(Dot)
@@ -10,3 +11,49 @@ def dot(x, y):
         return torch.matmul(x, y)
 
     return dot
+
+
+@pytorch_funcify.register(Max)
+def pytorch_funcify_Max(op, **kwargs):
+    axis = op.axis
+    keepdims = op.keepdims
+
+    def max(x):
+        return torch.max(x, dim=axis, keepdim=keepdims)
+
+    return max
+
+
+@pytorch_funcify.register(Argmax)
+def pytorch_funcify_Argmax(op, **kwargs):
+    axis = op.axis
+
+    def argmax(x):
+        if axis is None:
+            axes = tuple(range(x.ndim))
+        else:
+            axes = tuple(int(ax) for ax in axis)
+
+        # NumPy does not support multiple axes for argmax; this is a
+        # work-around
+        keep_axes = np.array([i for i in range(x.ndim) if i not in axes], dtype="int64")
+        # Not-reduced axes in front
+        transposed_x = np.transpose(
+            x, tuple(np.concatenate((keep_axes, np.array(axes, dtype="int64"))))
+        )
+        kept_shape = transposed_x.shape[: len(keep_axes)]
+        reduced_shape = transposed_x.shape[len(keep_axes) :]
+
+        # Numpy.prod returns 1.0 when arg is empty, so we cast it to int64
+        # Otherwise reshape would complain citing float arg
+        new_shape = (
+            *kept_shape,
+            np.prod(np.array(reduced_shape, dtype="int64"), dtype="int64"),
+        )
+        reshaped_x = torch.tensor(transposed_x.reshape(tuple(new_shape)))
+
+        max_idx_res = torch.argmax(reshaped_x, dim=-1).to(torch.long)
+
+        return max_idx_res
+
+    return argmax

tests/link/pytorch/test_math.py

@@ -1,8 +1,10 @@
 import numpy as np
+import pytest
 
 from pytensor.configdefaults import config
 from pytensor.graph.fg import FunctionGraph
 from pytensor.graph.op import get_test_value
+from pytensor.tensor.math import argmax, max
 from pytensor.tensor.type import matrix, scalar, vector
 from tests.link.pytorch.test_basic import compare_pytorch_and_py
 
@@ -28,3 +30,35 @@ def test_pytorch_dot():
     out = y.dot(alpha * A).dot(x) + beta * y
     fgraph = FunctionGraph([y, x, A, alpha, beta], [out])
     compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+
+@pytest.mark.parametrize(
+    "keepdims",
+    [True, False],
+)
+@pytest.mark.parametrize(
+    "axis",
+    [None, 1, 0],
+)
+def test_pytorch_max(axis, keepdims):
+    a = matrix("a", dtype=config.floatX)
+    a.tag.test_value = np.random.randn(4, 4).astype(config.floatX)
+    amx = max(a, axis=axis, keepdims=keepdims)
+    fgraph = FunctionGraph([a], amx, clone=False)
+    compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])
+
+
+@pytest.mark.parametrize(
+    "keepdims",
+    [True, False],
+)
+@pytest.mark.parametrize(
+    "axis",
+    [None, 1, (0,)],
+)
+def test_pytorch_argmax(axis, keepdims):
+    a = matrix("a", dtype=config.floatX)
+    a.tag.test_value = np.random.randn(4, 4).astype(config.floatX)
+    amx = argmax(a, axis=axis, keepdims=keepdims)
+    fgraph = FunctionGraph([a], amx)
+    compare_pytorch_and_py(fgraph, [get_test_value(i) for i in fgraph.inputs])