Commented code resolved

Author: Dhruvanshu-Joshi

pytensor/tensor/rewriting/uncanonicalize.py

@@ -34,33 +34,12 @@
 from pytensor.graph.rewriting.basic import copy_stack_trace, node_rewriter
 from pytensor.tensor.basic import Alloc, alloc, constant
 from pytensor.tensor.elemwise import DimShuffle
-
-# from pytensor.tensor.math import Argmax, Max, MaxAndArgmax, Min, neg
 from pytensor.tensor.math import Min, TensorMax, neg
 from pytensor.tensor.rewriting.basic import register_uncanonicalize
 from pytensor.tensor.shape import Reshape, reshape
 from pytensor.tensor.subtensor import Subtensor
 
 
-# @register_uncanonicalize
-# @node_rewriter([MaxAndArgmax])
-# def local_max_and_argmax(fgraph, node):
-#     """
-#     If we don't use the argmax, change it to a max only.
-#     """
-#     if isinstance(node.op, MaxAndArgmax):
-#         axis = node.op.axis
-#         if len(fgraph.clients[node.outputs[1]]) == 0:
-#             new = Max(axis)(node.inputs[0])
-#             copy_stack_trace(node.outputs[0], new)
-#             return [new, None]
-
-#         if len(fgraph.clients[node.outputs[0]]) == 0:
-#             new = Argmax(axis)(node.inputs[0])
-#             copy_stack_trace(node.outputs[0], new)
-#             return [None, new]
-
-
 @register_uncanonicalize
 @node_rewriter([neg])
 def local_max_to_min(fgraph, node):

tests/tensor/rewriting/test_uncanonicalize.py

@@ -22,69 +22,12 @@
 from tests.link.test_link import make_function
 
 
-# class TestMaxAndArgmax:
-#     def test_optimization(self):
-#         # If we use only the max output, we should replace this op with
-#         # a faster one.
-#         mode = pytensor.compile.mode.get_default_mode().including(
-#             "canonicalize", "fast_run"
-#         )
-
-#         for axis in [0, 1, -1]:
-#             n = matrix()
-
-#             f = function([n], max_and_argmax(n, axis)[0], mode=mode)
-#             topo = f.maker.fgraph.toposort()
-#             assert len(topo) == 1
-#             assert isinstance(topo[0].op, CAReduce)
-
-#             f = function([n], max_and_argmax(n, axis), mode=mode)
-#             topo = f.maker.fgraph.toposort()
-#             assert len(topo) == 1
-#             assert isinstance(topo[0].op, MaxAndArgmax)
-
-
 class TestMinMax:
     def setup_method(self):
         self.mode = pytensor.compile.mode.get_default_mode().including(
             "canonicalize", "fast_run"
         )
 
-    # def test_optimization_max(self):
-    #     data = np.asarray(np.random.random((2, 3)), dtype=config.floatX)
-    #     n = matrix()
-
-    #     for axis in [0, 1, -1]:
-    #         f = function([n], pt_max(n, axis), mode=self.mode)
-    #         topo = f.maker.fgraph.toposort()
-    #         assert len(topo) == 1
-    #         # assert isinstance(topo[0].op, CAReduce)
-    #         f(data)
-
-    #         f = function([n], pt_max(-n, axis), mode=self.mode)
-    #         topo = f.maker.fgraph.toposort()
-    #         import pytensor
-    #         pytensor.dprint(topo)
-    #         assert len(topo) == 2
-    #         assert isinstance(topo[0].op, Elemwise)
-    #         assert isinstance(topo[0].op.scalar_op, ps.Neg)
-    #         assert isinstance(topo[1].op, CAReduce)
-    #         f(data)
-
-    #         f = function([n], -pt_max(n, axis), mode=self.mode)
-    #         topo = f.maker.fgraph.toposort()
-    #         assert len(topo) == 2
-    #         assert isinstance(topo[0].op, CAReduce)
-    #         assert isinstance(topo[1].op, Elemwise)
-    #         assert isinstance(topo[1].op.scalar_op, ps.Neg)
-    #         f(data)
-
-    #         f = function([n], -pt_max(-n, axis), mode=self.mode)
-    #         topo = f.maker.fgraph.toposort()
-    #         assert len(topo) == 1
-    #         assert isinstance(topo[0].op, CAReduce)  # min
-    #         f(data)
-
     def test_optimization_min(self):
         data = np.asarray(np.random.random((2, 3)), dtype=config.floatX)
         n = matrix()