Added changes for seperating MaxandArgmax Op

Scalar problem solved

Finalise changes to seperate MaxAndArgmax Op

Author: Dhruvanshu-Joshi

pytensor/compile/function/__init__.py

@@ -312,6 +312,7 @@ def opt_log1p(node):
     else:
         # note: pfunc will also call orig_function -- orig_function is
         #      a choke point that all compilation must pass through
+
         fn = pfunc(
             params=inputs,
             outputs=outputs,

pytensor/compile/function/types.py

@@ -1758,6 +1758,7 @@ def orig_function(
             name=name,
             fgraph=fgraph,
         )
+        print(m)
         with config.change_flags(compute_test_value="off"):
             fn = m.create(defaults)
     finally:

pytensor/graph/op.py

@@ -291,6 +291,7 @@ def __call__(
 
         """
         node = self.make_node(*inputs, **kwargs)
+
         if name is not None:
             if len(node.outputs) == 1:
                 node.outputs[0].name = name

pytensor/ifelse.py

@@ -477,7 +477,7 @@ def cond_make_inplace(fgraph, node):
     Reshape,
     Unbroadcast,
     pt.math.Dot,
-    pt.math.TensorMax,
+    pt.math.Max,
     pt.math.Argmax,
     pt.subtensor.Subtensor,
     pt.subtensor.IncSubtensor,

pytensor/link/jax/dispatch/nlinalg.py

@@ -2,7 +2,7 @@
 
 from pytensor.link.jax.dispatch import jax_funcify
 from pytensor.tensor.blas import BatchedDot
-from pytensor.tensor.math import Dot, MaxAndArgmax
+from pytensor.tensor.math import Argmax, Dot, Max
 from pytensor.tensor.nlinalg import (
     SVD,
     Det,
@@ -104,18 +104,73 @@ def batched_dot(a, b):
     return batched_dot
 
 
-@jax_funcify.register(MaxAndArgmax)
-def jax_funcify_MaxAndArgmax(op, **kwargs):
+# @jax_funcify.register(Max)
+# @jax_funcify.register(Argmax)
+# def jax_funcify_MaxAndArgmax(op, **kwargs):
+#     axis = op.axis
+
+#     def maxandargmax(x, axis=axis):
+#         if axis is None:
+#             axes = tuple(range(x.ndim))
+#         else:
+#             axes = tuple(int(ax) for ax in axis)
+
+#         max_res = jnp.max(x, axis)
+
+#         # NumPy does not support multiple axes for argmax; this is a
+#         # work-around
+#         keep_axes = jnp.array(
+#             [i for i in range(x.ndim) if i not in axes], dtype="int64"
+#         )
+#         # Not-reduced axes in front
+#         transposed_x = jnp.transpose(
+#             x, jnp.concatenate((keep_axes, jnp.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,
+#             jnp.prod(jnp.array(reduced_shape, dtype="int64"), dtype="int64"),
+#         )
+#         reshaped_x = transposed_x.reshape(new_shape)
+
+#         max_idx_res = jnp.argmax(reshaped_x, axis=-1).astype("int64")
+
+#         return max_res, max_idx_res
+
+#     return maxandargmax
+
+
+@jax_funcify.register(Max)
+def jax_funcify_Max(op, **kwargs):
     axis = op.axis
 
-    def maxandargmax(x, axis=axis):
+    def max(x, axis=axis):
+        # if axis is None:
+        #     axes = tuple(range(x.ndim))
+        # else:
+        #     axes = tuple(int(ax) for ax in axis)
+
+        max_res = jnp.max(x, axis)
+
+        return max_res
+
+    return max
+
+
+@jax_funcify.register(Argmax)
+def jax_funcify_Argmax(op, **kwargs):
+    axis = op.axis
+
+    def argmax(x, axis=axis):
         if axis is None:
             axes = tuple(range(x.ndim))
         else:
             axes = tuple(int(ax) for ax in axis)
 
-        max_res = jnp.max(x, axis)
-
         # NumPy does not support multiple axes for argmax; this is a
         # work-around
         keep_axes = jnp.array(
@@ -138,6 +193,6 @@ def maxandargmax(x, axis=axis):
 
         max_idx_res = jnp.argmax(reshaped_x, axis=-1).astype("int64")
 
-        return max_res, max_idx_res
+        return max_idx_res
 
-    return maxandargmax
+    return argmax

pytensor/link/numba/dispatch/elemwise.py

@@ -44,7 +44,7 @@
 )
 from pytensor.scalar.basic import add as add_as
 from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
-from pytensor.tensor.math import MaxAndArgmax, MulWithoutZeros, Sum
+from pytensor.tensor.math import Argmax, MulWithoutZeros, Sum
 from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
 from pytensor.tensor.type import scalar
 
@@ -985,8 +985,8 @@ def log_softmax_py_fn(x):
     return log_softmax
 
 
-@numba_funcify.register(MaxAndArgmax)
-def numba_funcify_MaxAndArgmax(op, node, **kwargs):
+@numba_funcify.register(Argmax)
+def numba_funcify_Argmax(op, node, **kwargs):
     axis = op.axis
     x_at = node.inputs[0]
     x_dtype = x_at.type.numpy_dtype
@@ -996,8 +996,8 @@ def numba_funcify_MaxAndArgmax(op, node, **kwargs):
     if x_ndim == 0:
 
         @numba_basic.numba_njit(inline="always")
-        def maxandargmax(x):
-            return x, 0
+        def argmax(x):
+            return 0
 
     else:
         axes = tuple(int(ax) for ax in axis)
@@ -1006,20 +1006,6 @@ def maxandargmax(x):
         # work-around
         keep_axes = tuple(i for i in range(x_ndim) if i not in axes)
 
-        reduce_max_py_fn = create_multiaxis_reducer(
-            scalar_maximum,
-            -np.inf,
-            axes,
-            x_ndim,
-            x_dtype,
-            return_scalar=False,
-        )
-        reduce_max = jit_compile_reducer(
-            Apply(node.op, node.inputs, [node.outputs[0].clone()]),
-            reduce_max_py_fn,
-            reduce_to_scalar=False,
-        )
-
         reduced_x_ndim = x_ndim - len(axes) + 1
         argmax_axis = create_axis_apply_fn(
             np.argmax, reduced_x_ndim - 1, reduced_x_ndim, np.int64
@@ -1030,9 +1016,7 @@ def maxandargmax(x):
         sl2 = slice(len(keep_axes), None)
 
         @numba_basic.numba_njit
-        def maxandargmax(x):
-            max_res = reduce_max(x)
-
+        def argmax(x):
             # Not-reduced axes in front
             transposed_x = np.ascontiguousarray(np.transpose(x, reaxis_order))
             kept_shape = transposed_x.shape[sl1]
@@ -1048,6 +1032,6 @@ def maxandargmax(x):
 
             max_idx_res = argmax_axis(reshaped_x)
 
-            return max_res, max_idx_res
+            return max_idx_res
 
-    return maxandargmax
+    return argmax