code: 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)
  
imports: import numpy
  import pytensor
  import pytensor.tensor
  from pytensor import function
  from pytensor import scalar
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.rewriting.basic import out2in
  from pytensor.link.basic import PerformLinker
  from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
  from pytensor.tensor.math import MaxAndArgmax, max_and_argmax
  from pytensor.tensor.math import max
  from pytensor.tensor.math import min
  from pytensor.tensor.rewriting.uncanonicalize import local_alloc_dimshuffle, local_dimshuffle_alloc, local_dimshuffle_subtensor, local_reshape_dimshuffle
  from pytensor.tensor.shape import reshape, specify_shape
  from pytensor.tensor.type import dtensor4, iscalar, matrix, tensor, vector
  from tests.link.test_link import make_function
label_desc: Use image node_id 1 for calling the TestMaxAndArgmax obj's underlying member method code with example usage: obj.test_optimization() without return types