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

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

tests/tensor/test_math.py

pytensor/link/numba/dispatch/elemwise.py

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 81.69014% with 13 lines in your changes missing coverage. Please review.

Project coverage is 80.87%. Comparing base (15b90be) to head (d953a0d).
Report is 227 commits behind head on main.

Files with missing lines	Patch %	Lines
pytensor/tensor/math.py	78.57%	6 Missing and 6 partials ⚠️
pytensor/link/jax/dispatch/nlinalg.py	92.30%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #731      +/-   ##
==========================================
+ Coverage   80.83%   80.87%   +0.03%     
==========================================
  Files         162      163       +1     
  Lines       46862    46847      -15     
  Branches    11465    11463       -2     
==========================================
+ Hits        37881    37887       +6     
- Misses       6733     6747      +14     
+ Partials     2248     2213      -35

Files with missing lines	Coverage Δ
pytensor/ifelse.py	`51.70% <ø> (ø)`
pytensor/link/numba/dispatch/elemwise.py	`91.85% <100.00%> (+3.13%)`	⬆️
pytensor/tensor/rewriting/uncanonicalize.py	`96.63% <100.00%> (+0.42%)`	⬆️
pytensor/link/jax/dispatch/nlinalg.py	`90.47% <92.30%> (+0.73%)`	⬆️
pytensor/tensor/math.py	`90.43% <78.57%> (+0.77%)`	⬆️

... and 37 files with indirect coverage changes

ricardoV94

Almost there, just some cleanup needed

pytensor/compile/function/types.py

pytensor/link/jax/dispatch/nlinalg.py

pytensor/tensor/math.py

pytensor/tensor/rewriting/uncanonicalize.py

tests/link/numba/test_basic.py

ricardoV94

Small notes. Reverting the white space changes is useful, because git diff will then only show the files that have meaningful changes from this PR

pytensor/graph/op.py

pytensor/compile/function/__init__.py

pytensor/ifelse.py

pytensor/link/jax/dispatch/nlinalg.py

ricardoV94 · 2024-06-10T14:26:39Z

pytensor/tensor/math.py

+def __getattr__(name):
+    if name == "MaxandArgmax":
+        warnings.warn(
+            "The class `MaxandArgmax` has been deprecated. "
+            "Call `Max` and `Argmax` seperately as an alternative.",
+            FutureWarning,
+        )
+


This has to raise some sort of error because MaxAndArgmax no longer exist. It can't simply be a warning. I think it's better to just remove it

should we raise an AttributeError here?

Whatever error would be raised if this helper was not here, AttrubetError sounds about right but confirm with trying without the special code

The helper also needs to raise the standard error for things other than MaxandArgmax

Confirmed that AttributeError is raised

This is how we implement it in pymc for reference: https://github.com/pymc-devs/pymc/blob/f44071bdda363f743548187d3e124a027adfdb77/pymc/distributions/transforms.py#L54-L67

Although according to the PEP it should probably have a !r in the standard AttributeError message: https://peps.python.org/pep-0562/#rationale

In our case, we won't have to give a warning right?

Yes in our case we just raise a more informative AttributeError instead of the default one

In your last change you made this useless, it's the same as if we didn't implement __getattr__ at all. The point of adding it is to have a custom message for the deprecated Op, but have the default behavior for everything else.

Rectified this

tests/tensor/test_math.py

tests/tensor/utils.py

pytensor/tensor/math.py

ricardoV94 · 2024-06-13T08:51:05Z

Almost there, looks great. Just 3 unresolved comments above

ricardoV94 · 2024-06-13T09:25:01Z

tests/tensor/test_math.py

@@ -1404,6 +1422,11 @@ def test_bool(self):
        assert np.all(i)


+def test_MaxAndArgmax_deprecated():
+    with pytest.raises(AttributeError):


Add the specific message we expect for trying to access the MaxAndArgmax

Suggested change

with pytest.raises(AttributeError):

with pytest.raises(AttributeError, match=...):

ricardoV94

Awesome!

ricardoV94 · 2024-06-13T09:53:52Z

Some tests failing, it seems there are still some internal imports of MaxAndArgmax

Dhruvanshu-Joshi · 2024-06-13T10:17:49Z

Some tests failing, it seems there are still some internal imports of MaxAndArgmax

My bad. I rectified this. JAX tests are skipped locally so could not test them anywhere other than the CI. Hopefully, they'll pass now.

ricardoV94 · 2024-06-13T12:54:12Z

Thanks @Dhruvanshu-Joshi !