Deprecate and raise AttributeError for MaxAndArgmax

Author: Dhruvanshu-Joshi

pytensor/tensor/math.py

@@ -110,9 +110,9 @@
 def __getattr__(name):
     if name == "MaxAndArgmax":
         raise AttributeError(
-            "The class `MaxAndArgmax` has been deprecated. "
-            "Call `Max` and `Argmax` separately as an alternative."
+            "The class `MaxandArgmax` has been deprecated. Call `Max` and `Argmax` seperately as an alternative."
         )
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
 
 
 def _get_atol_rtol(a, b):
@@ -565,16 +565,6 @@ def max(x, axis=None, keepdims=False):
     We return an error as numpy when we reduce a dim with a shape of 0.
 
     """
-
-    # We have a choice of implementing this call with the
-    # CAReduce op or the MaxAndArgmax op.
-
-    # MaxAndArgmax supports grad and Rop, so we prefer to use that.
-    # CAReduce is faster, but optimizations will replace MaxAndArgmax[0]
-    # with CAReduce at compile time, so at this stage the important
-    # thing is supporting all user interface features, not speed.
-    # Some cases can be implemented only with CAReduce.
-
     out = max_and_argmax(x, axis)[0]
 
     if keepdims:

pytensor/tensor/rewriting/uncanonicalize.py

@@ -52,7 +52,7 @@ def local_max_to_min(fgraph, node):
     Notes
     -----
     We don't need an opt that will do the reverse as by default
-    the interface put only MaxAndArgmax into the graph.
+    the interface put only Max into the graph.
 
     """
     if node.op == neg and node.inputs[0].owner:

tests/link/jax/test_nlinalg.py

@@ -11,7 +11,7 @@
 from pytensor.link.jax import JAXLinker
 from pytensor.tensor import blas as pt_blas
 from pytensor.tensor import nlinalg as pt_nlinalg
-from pytensor.tensor.math import MaxAndArgmax, maximum
+from pytensor.tensor.math import Argmax, Max, maximum
 from pytensor.tensor.math import max as pt_max
 from pytensor.tensor.type import dvector, matrix, scalar, tensor3, vector
 from tests.link.jax.test_basic import compare_jax_and_py
@@ -88,7 +88,8 @@ def test_jax_basic_multiout_omni():
     # Test that a single output of a multi-output `Op` can be used as input to
     # another `Op`
     x = dvector()
-    mx, amx = MaxAndArgmax([0])(x)
+    mx = Max([0])(x)
+    amx = Argmax([0])(x)
     out = mx * amx
     out_fg = FunctionGraph([x], [out])
     compare_jax_and_py(out_fg, [np.r_[1, 2]])

tests/tensor/test_math.py

@@ -1032,7 +1032,6 @@ def test_vectorize(self, core_axis, batch_axis):
         x = tensor(shape=(5, 5, 5, 5))
         batch_x = tensor(shape=(3, 5, 5, 5, 5))
 
-        # Test MaxAndArgmax
         argmax_x = argmax(x, axis=core_axis)
 
         arg_max_node = argmax_x.owner
@@ -1423,7 +1422,10 @@ def test_bool(self):
 
 
 def test_MaxAndArgmax_deprecated():
-    with pytest.raises(AttributeError):
+    with pytest.raises(
+        AttributeError,
+        match="The class `MaxandArgmax` has been deprecated. Call `Max` and `Argmax` seperately as an alternative.",
+    ):
         pytensor.tensor.math.MaxAndArgmax