Finalise chnages to seperate MaxAndArgmax Op

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -28,6 +28,7 @@
     constant,
     stack,
     switch,
+    zeros_like,
 )
 from pytensor.tensor.blockwise import Blockwise, vectorize_node_fallback
 from pytensor.tensor.elemwise import (
@@ -192,10 +193,9 @@ def perform(self, node, inp, outs):
         (x,) = inp
         axes = self.axis
         (max_idx,) = outs
+
         if axes is None:
             axes = tuple(range(x.ndim))
-        # else:
-        #     axes = tuple(int(ax) for ax in axes)
 
         # Numpy does not support multiple axes for argmax
         # Work around
@@ -400,14 +400,22 @@ def max_and_argmax(a, axis=None, keepdims=False):
     # Check axis and convert it to a Python list of integers.
     # Axis will be used as an op param of MaxAndArgmax.
     a = as_tensor_variable(a)
+
+    flag = False
+    if axis == ():
+        flag = True
+
     axis = check_and_normalize_axes(a, axis)
-    if len(axis) == 0:
-        # axis = list(range(a.type.ndim))
+
+    if len(axis) == 0 and not flag:
         axis = None
-    # out = TensorMax(axis)(a)
+
     out = Max(axis)(a)
-    argout = Argmax(axis)(a)
-    # _, argout = MaxAndArgmax(axis)(a)
+
+    if not flag:
+        argout = Argmax(axis)(a)
+    else:
+        argout = zeros_like(a, dtype="int64")
 
     if keepdims:
         out = makeKeepDims(a, out, axis)
@@ -483,7 +491,7 @@ def grad(self, inp, grads):
 
         g_max_disconnected = isinstance(g_max.type, DisconnectedType)
 
-        # if the op is tota lly disconnected, so are its inputs
+        # if the op is totally disconnected, so are its inputs
         if g_max_disconnected:
             return [DisconnectedType()()]
 

tests/link/numba/test_elemwise.py

@@ -552,24 +552,6 @@ def test_LogSoftmax(x, axis, exc):
         ),
     ],
 )
-# def test_MaxAndArgmax(x, axes, exc):
-#     g = ptm.MaxAndArgmax(axes)(x)
-
-#     if isinstance(g, list):
-#         g_fg = FunctionGraph(outputs=g)
-#     else:
-#         g_fg = FunctionGraph(outputs=[g])
-
-#     cm = contextlib.suppress() if exc is None else pytest.warns(exc)
-#     with cm:
-#         compare_numba_and_py(
-#             g_fg,
-#             [
-#                 i.tag.test_value
-#                 for i in g_fg.inputs
-#                 if not isinstance(i, SharedVariable | Constant)
-#             ],
-#         )
 def test_Max(x, axes, exc):
     g = ptm.Max(axes)(x)
 

tests/tensor/test_math.py

@@ -764,8 +764,7 @@ def setup_method(self):
         Max.debug = 0
         Argmax.debug = 0
 
-    def test_basic_0(self):
-        # dbt: for some reason, Argmax does not work when I pass: n = as_tensor_variable(5.0)
+    def test_basic(self):
         n = as_tensor_variable(5)
         v, i = eval_outputs(max_and_argmax(n, axis=()))
         assert v == 5.0
@@ -1040,6 +1039,29 @@ def test_vectorize(self, core_axis, batch_axis):
         assert isinstance(new_node.op, Argmax)
         assert new_node.op.axis == batch_axis
 
+    def test_max_empty_axis(self):
+        x = np.random.normal(size=(2, 3, 5, 7))
+        axis = ()
+
+        non_axis = tuple(i for i in range(x.ndim) if i not in axis)
+        shape_axis = tuple(x.shape[dim] for dim in axis)
+        shape_non_axis = tuple(x.shape[dim] for dim in non_axis)
+        x_transposed = x.transpose(*axis, *non_axis)
+
+        x_axis_raveled = x_transposed.reshape(
+            np.prod(shape_axis, dtype=int), np.prod(shape_non_axis, dtype=int)
+        )
+        max_x = max_and_argmax(x, axis=axis)[0].eval()
+        argmax_x = max_and_argmax(x, axis=axis)[1].eval()
+
+        raveled_max = x_axis_raveled[
+            argmax_x.ravel(), np.arange(np.prod(shape_non_axis, dtype=int))
+        ]
+        indirect_max = raveled_max.reshape(shape_non_axis)
+
+        np.testing.assert_allclose(max_x, x.max(axis=axis))
+        np.testing.assert_allclose(indirect_max, x.max(axis=axis))
+
 
 class TestArgminArgmax:
     def setup_method(self):
@@ -1379,19 +1401,11 @@ def test_uint(self):
             data = np.array([itype.min + 3, itype.min, itype.max - 5, itype.max], dtype)
             n = as_tensor_variable(data)
             assert min(n).dtype == dtype
-            # print(min(n).owner.inputs[1].acc_dtype)
             i = eval_outputs(min(n))
-            # pytensor.dprint(n)
-            # for x in n:
-            #     print(x.eval())
-            print(i)
-            print(itype.min)
-            print()
             assert i == itype.min
-            # assert max(n).dtype == dtype
-            # i = eval_outputs(max(n))
-            # assert i == itype.max
-            # assert 0
+            assert max(n).dtype == dtype
+            i = eval_outputs(max(n))
+            assert i == itype.max
 
     def test_bool(self):
         data = np.array([True, False], "bool")