Intermediate changes

Author: Dhruvanshu-Joshi

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, Max, MulWithoutZeros, Sum
 from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
 from pytensor.tensor.type import scalar
 
@@ -985,8 +985,78 @@ def log_softmax_py_fn(x):
     return log_softmax
 
 
-@numba_funcify.register(MaxAndArgmax)
-def numba_funcify_MaxAndArgmax(op, node, **kwargs):
+# @numba_funcify.register(Max)
+# @numba_funcify.register(Argmax)
+# # @numba_funcify.register(MaxandArgmax)
+# def numba_funcify_MaxAndArgmax(op, node, **kwargs):
+#     axis = op.axis
+#     x_at = node.inputs[0]
+#     x_dtype = x_at.type.numpy_dtype
+#     x_dtype = numba.np.numpy_support.from_dtype(x_dtype)
+#     x_ndim = x_at.ndim
+
+#     if x_ndim == 0:
+
+#         @numba_basic.numba_njit(inline="always")
+#         def maxandargmax(x):
+#             return x, 0
+
+#     else:
+#         axes = tuple(int(ax) for ax in axis)
+
+#         # NumPy does not support multiple axes for argmax; this is a
+#         # 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
+#         )
+
+#         reaxis_order = keep_axes + axes
+#         sl1 = slice(None, len(keep_axes))
+#         sl2 = slice(len(keep_axes), None)
+
+#         @numba_basic.numba_njit
+#         def maxandargmax(x):
+#             max_res = reduce_max(x)
+
+#             # Not-reduced axes in front
+#             transposed_x = np.ascontiguousarray(np.transpose(x, reaxis_order))
+#             kept_shape = transposed_x.shape[sl1]
+#             reduced_shape = transposed_x.shape[sl2]
+#             reduced_size = 1
+#             for s in reduced_shape:
+#                 reduced_size *= s
+
+#             # 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, reduced_size)
+#             reshaped_x = transposed_x.reshape(new_shape)
+
+#             max_idx_res = argmax_axis(reshaped_x)
+
+#             return max_res, max_idx_res
+
+#     return maxandargmax
+
+
+@numba_funcify.register(Max)
+def numba_funcify_Max(op, node, **kwargs):
     axis = op.axis
     x_at = node.inputs[0]
     x_dtype = x_at.type.numpy_dtype
@@ -996,15 +1066,15 @@ def numba_funcify_MaxAndArgmax(op, node, **kwargs):
     if x_ndim == 0:
 
         @numba_basic.numba_njit(inline="always")
-        def maxandargmax(x):
-            return x, 0
+        def max(x):
+            return x
 
     else:
         axes = tuple(int(ax) for ax in axis)
 
         # NumPy does not support multiple axes for argmax; this is a
         # work-around
-        keep_axes = tuple(i for i in range(x_ndim) if i not in axes)
+        # keep_axes = tuple(i for i in range(x_ndim) if i not in axes)
 
         reduce_max_py_fn = create_multiaxis_reducer(
             scalar_maximum,
@@ -1020,6 +1090,50 @@ def maxandargmax(x):
             reduce_to_scalar=False,
         )
 
+        @numba_basic.numba_njit
+        def max(x):
+            max_res = reduce_max(x)
+
+            return max_res
+
+    return max
+
+
+@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
+    x_dtype = numba.np.numpy_support.from_dtype(x_dtype)
+    x_ndim = x_at.ndim
+
+    if x_ndim == 0:
+
+        @numba_basic.numba_njit(inline="always")
+        def argmax(x):
+            return 0
+
+    else:
+        axes = tuple(int(ax) for ax in axis)
+
+        # NumPy does not support multiple axes for argmax; this is a
+        # 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 +1144,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 +1160,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