Break MaxandArgmax Op to seperate TensorMax Op and Argmax Op #731

Dhruvanshu-Joshi · 2024-04-27T02:18:43Z

Description

MaxandArgmax Op calculates both maximum and argmax together. With this PR, we aim to have seperate ops for the two operations.

Related Issue

Closes Remove MaxAndArgmax Op #334
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

pytensor/tensor/math.py

pytensor/tensor/rewriting/uncanonicalize.py

tests/tensor/rewriting/test_uncanonicalize.py

Dhruvanshu-Joshi · 2024-04-27T02:26:58Z

tests/tensor/test_max_argmax.py

+        TensorMax.debug = 0
+
+    def test_basic(self):
+        # dbt: for some reason, Argmax does not work when I pass: n = as_tensor_variable(5.0)


For some reason, Argmax does not work when I pass: n = as_tensor_variable(5.0). MaxandArgmax used to work fine.

Hmm, what does numpy do with scalar arrays? We should do the same as them

numpy seems to work fine

The/a problem seems to be this line, which is forcing Axis to be a tuple, when it should be allowed to be None for the scalar case:

pytensor/pytensor/tensor/math.py

Line 362 in 27bd9aa

self.axis = tuple(axis)

Should be self.axis = axis

Tried this and it does not seem to work.
This is the line in particular that causes a problem:

pytensor/pytensor/link/c/basic.py

Lines 1745 to 1774 in e48ff56

def __call__(self):

failure = self.run_cthunk(self.cthunk)

if failure:

task, taskname, id = self.find_task(failure)

try:

trace = task.trace

except AttributeError:

trace = ()

try:

exc_type, _exc_value, exc_trace = self.error_storage

if task in self.nodes:

self.position_of_error = self.nodes.index(task)

# this can be used to retrieve the location the Op was declared

exc_value = exc_type(_exc_value)

exc_value.__thunk_trace__ = trace

except Exception:

print(

(

"ERROR retrieving error_storage."

"Was the error set in the c code?"

),

end=" ",

file=sys.stderr,

)

print(self.error_storage, file=sys.stderr)

raise

raise exc_value.with_traceback(exc_trace)

def __str__(self):

return f"{type(self).__name__}({self.module})"

The if failure block never executes in case of non scalar inputs.

pytensor/pytensor/link/c/cutils.py

Lines 102 to 103 in e48ff56

try:

from cutils_ext.cutils_ext import * # noqa

run_cthunk is imported using this but I cannot find cutils_ext.

You may be hitting a windows/installation issue then. Does the test fail in the CI here as well?

Yes the CI error is the same as what I face locally. The difference between a scalar like as_tensor_variable(5.0) and non-scalar like as_tensor_variable([5.0]) is that in the former case, the if failure: block executes and in case of non-scalars, it never does.

tests/tensor/test_max_argmax.py

pytensor/tensor/math.py

pytensor/tensor/rewriting/uncanonicalize.py

tests/tensor/rewriting/test_uncanonicalize.py

tests/tensor/test_max_argmax.py

pytensor/tensor/math.py

Dhruvanshu-Joshi · 2024-04-29T08:09:20Z

tests/tensor/test_math.py


    def test_basic(self):
+        # dbt: for some reason, Argmax does not work when I pass: n = as_tensor_variable(5.0)


Scalars are still a problem. Am working on this and looking for how numpy handles them as suggested.

It's raising from this C-code check:

pytensor/pytensor/tensor/math.py

Line 227 in 8e7f626

"Argmax, bad axis argument");

It seems not to handle the empty axes case that's passed when it's a scalar. Instead of empty axes we can convert to None or zero which is equivalent for the scalar. I prefer None because that's the default anyway.

But actually these lines seem to be creating the problem in the first place?

pytensor/pytensor/tensor/math.py

Lines 403 to 405 in 8e7f626

axis = check_and_normalize_axes(a, axis)

if len(axis) == 0:

axis = list(range(a.type.ndim))

I don't think it should be needed or this convoluted (referring to that check_and_normalize_axes, which I think is only used here?). We handle axes in other places with way less code. We should use the numpy helper like we do for other cases, or let axes = None alone, which both argmax and max support anyway.

I was not seeing problems before because I was creating Argmax(axis=None)(pt.as_tensor(5.0)) directly. Which shows the problem is how the helper is creating the Argmax, basically Argmax(axis=()). Actually I am not sure what Argmax(axis=()) should do, I think it should return zeros_like(x) since it corresponds to a no-reduction. np.max(x, axis=()) just returns x as well. We should check our Max does the same btw, which I think it's not doing.

I tried the

It's raising from this C-code check:

pytensor/pytensor/tensor/math.py

Line 227 in 8e7f626

"Argmax, bad axis argument");

It seems not to handle the empty axes case that's passed when it's a scalar. Instead of empty axes we can convert to None or zero which is equivalent for the scalar. I prefer None because that's the default anyway.

But actually these lines seem to be creating the problem in the first place?

pytensor/pytensor/tensor/math.py

Lines 403 to 405 in 8e7f626

axis = check_and_normalize_axes(a, axis)

if len(axis) == 0:

axis = list(range(a.type.ndim))

I don't think it should be needed or this convoluted (referring to that check_and_normalize_axes, which I think is only used here?). We handle axes in other places with way less code. We should use the numpy helper like we do for other cases, or let axes = None alone, which both argmax and max support anyway.

I tried this locally and it works with a silly modification:

def max_and_argmax(a, axis=None, keepdims=False): """ Returns maximum elements and their indices obtained by iterating over given axis. When axis is None (the default value), the max is performed over the flattened tensor. Parameters ---------- keepdims : bool If this is set to True, the axes which are reduced are left in the result as dimensions with size one. With this option, the result will broadcast correctly against the original tensor. """ # 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) axis = check_and_normalize_axes(a, axis) if len(axis) == 0: axis = None out = Max(axis)(a) argout = Argmax(axis)(a) if keepdims: out = makeKeepDims(a, out, axis) argout = makeKeepDims(a, argout, axis) return [out, argout]

Scalars work this way. But now in the grad function for max, the line axis = as_tensor_variable(self.axis) misbehaves as self.axis is None and as_tensor_variable(None) is wrong. Maybe doing if self.axis is None: self.axis= tuple(range(x.ndim)) help but will it be correct?

I was not seeing problems before because I was creating Argmax(axis=None)(pt.as_tensor(5.0)) directly. Which shows the problem is how the helper is creating the Argmax, basically Argmax(axis=()). Actually I am not sure what Argmax(axis=()) should do, I think it should return zeros_like(x) since it corresponds to a no-reduction. np.max(x, axis=()) just returns x as well. We should check our Max does the same btw, which I think it's not doing.

Do the current changes reflect on this? The case of pt.as_tensor(5.0) is handled effectively now ig.
And the assert

pytensor/tests/tensor/test_math.py

Line 768 in e48ff56

assert v == 5.0

does not give any error so I assume it is doing what we expect it to do?

Even

n = as_tensor_variable(5.0) v, i = eval_outputs(max_and_argmax(n, axis=())) assert v == 5.0 assert i == 0 assert i.dtype == "int64" v = eval_outputs(max_and_argmax(n)[0].shape) assert len(v) == 0 v = eval_outputs(max_and_argmax(n)[1].shape) assert len(v) == 0

works fine so I assume axis=() works as expected?

I don't know why grad was converting axis to a tensor_variable, there is no point, since axis have to be constant. You can avoid that conversion.

Dhruvanshu-Joshi · 2024-04-29T08:16:33Z

tests/tensor/test_math.py

@@ -1386,6 +1383,12 @@ def test_uint(self):
            n = as_tensor_variable(data)
            assert min(n).dtype == dtype
            i = eval_outputs(min(n))
+            # pytensor.dprint(n)


The dtype uint64 fails strangely for some reason. The rest work fine. The error message is that itype.min = 0 but i comes out to be equal to the 18446744073709551610 which is the second maximum in the list of [0 , 3, 18446744073709551610, 18446744073709551615].

The error is:

> assert i == itype.min E assert array(18446744073709551610, dtype=uint64) == 0 E + where 0 = iinfo(min=0, max=18446744073709551615, dtype=uint64).min

Might have to do with the dtype used for internal accumulation.

Can you elaborate a little on this?

I am not sure what the problem is yet without looking further. But CAReduce has 2 dtypes, the output one and the one used for internal accumulation. I was wondering if the problem was coming from the internal accumulation dtype. Also uint are tricky because they don't represent negative numbers, but I think our implementation of min is something like -max(-x). You may need to investigate a bit the behavior to understand what's going on.

Dhruvanshu-Joshi · 2024-05-01T14:48:54Z

pytensor/link/numba/dispatch/elemwise.py

-        def maxandargmax(x):
-            return x, 0
+        def max(x):
+            return x

    else:
        axes = tuple(int(ax) for ax in axis)


This causes problem when axis is None.

You'll need to convert into axis = tuple(range(x_ndim)). I assume this is only a problem for Argmax? I think the conversion is already done for Max by default (as is for all CAReduce)?

We can do the same conversion for Argmax. Are we always creating Argmax for the user in pt.argmax? If so we can do the conversion there already

Yeh this code was redundant. numba_funcify_CAReduce already handles it effectively. Argmax behaves as expected already.

pytensor/tensor/math.py

Scalar problem solved Finalise changes to seperate MaxAndArgmax Op

Dhruvanshu-Joshi · 2024-05-16T10:21:22Z

The tests failing are because of uint64 data type which is highlighted in #770 . So for this to be ready, should I just remove the test for uint64 for now and open another issue to add support back for this test once #770 is solved?

ricardoV94 · 2024-05-16T12:36:11Z

The tests failing are because of uint64 data type which is highlighted in #770 . So for this to be ready, should I just remove the test for uint64 for now and open another issue to add support back for this test once #770 is solved?

You can mark the test with pytest.mark.xfail. There are a couple of examples in the codebase

codecov · 2024-05-22T16:55:55Z

Codecov Report

Attention: Patch coverage is 79.22078% with 16 lines in your changes are missing coverage. Please review.

Project coverage is 80.85%. Comparing base (15b90be) to head (6b07a6e).
Report is 11 commits behind head on main.

Additional details and impacted files