Scalar problem solved

Author: Dhruvanshu-Joshi

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 Argmax, Max, 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,120 +985,6 @@ def log_softmax_py_fn(x):
     return log_softmax
 
 
-# @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
-    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 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)
-
-        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,
-        )
-
-        @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
@@ -1120,20 +1006,6 @@ def argmax(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

pytensor/tensor/math.py

@@ -194,8 +194,8 @@ def perform(self, node, inp, outs):
         (max_idx,) = outs
         if axes is None:
             axes = tuple(range(x.ndim))
-        else:
-            axes = tuple(int(ax) for ax in axes)
+        # else:
+        #     axes = tuple(int(ax) for ax in axes)
 
         # Numpy does not support multiple axes for argmax
         # Work around
@@ -402,7 +402,8 @@ def max_and_argmax(a, axis=None, keepdims=False):
     a = as_tensor_variable(a)
     axis = check_and_normalize_axes(a, axis)
     if len(axis) == 0:
-        axis = list(range(a.type.ndim))
+        # axis = list(range(a.type.ndim))
+        axis = None
     # out = TensorMax(axis)(a)
     out = Max(axis)(a)
     argout = Argmax(axis)(a)
@@ -475,6 +476,8 @@ def grad(self, inp, grads):
         # g_max to x's shape when axis=0 the broadcasting mechanism
         # does it automatically
         x = inp[0]
+        if self.axis is None:
+            self.axis = tuple(range(x.ndim))
         axis = as_tensor_variable(self.axis)
         (g_max,) = grads
 
@@ -617,12 +620,6 @@ def min(x, axis=None, keepdims=False):
     elif str_x_type in uint_dtypes:
         itype = np.iinfo(x.dtype)
         max_val = np.array(itype.max, dtype=itype.dtype)
-        # print('a')
-        # for c in (max_val - x):
-        #     print(c.eval())
-        # print()
-        # print(max(max_val - x, axis=axis, keepdims=keepdims).eval())
-        # print()
         return max_val - max(max_val - x, axis=axis, keepdims=keepdims)
     elif str_x_type == "bool":
         return ~max(~x, axis=axis, keepdims=keepdims)

tests/link/numba/test_basic.py

@@ -256,6 +256,8 @@ def compare_numba_and_py(
     if assert_fn is None:
 
         def assert_fn(x, y):
+            print(x)
+            print(y)
             return np.testing.assert_allclose(x, y, rtol=1e-4) and compare_shape_dtype(
                 x, y
             )