class_docstr: Currently we only use it to infer caffe shape, we use tensorflow engine
  to infer tensorflow op shapes, since tensorflow has too many ops.
class_docstr_tok: ["Currently","we","only","use","it","to","infer","caffe","shape",",","we","use","tensorflow","engine","to","infer","tensorflow","op","shapes",",","since","tensorflow","has","too","many","ops","."]
code: def infer_shape_argmax(self, op):
      input_shape = self._output_shape_cache[op.input[0]]
      output_dim_num = len(input_shape)
      if output_dim_num < 3:
          output_dim_num = 3
  
      axis_arg = ConverterUtil.get_arg(op, MaceKeyword.mace_axis_str)
      has_axis = axis_arg is not None
      axis_value = 0
      if has_axis:
          axis_value = axis_arg.i
          if axis_value < 0:
              axis_value = len(input_shape) + axis_value
  
      top_k = ConverterUtil.get_arg(op, MaceKeyword.mace_top_k_str).i
      mace_check(top_k >= 1, "Invalid top_k value")
      out_val = ConverterUtil.get_arg(
          op, MaceKeyword.mace_out_val_str
      ).i
  
      if has_axis:  # Produces max_ind or max_val per axis
          output_shape = input_shape
          output_shape[axis_value] = top_k
      else:
          output_shape = [1] * output_dim_num
          output_shape[0] = input_shape[0]
          output_shape[2] = top_k
          if out_val:  # Produces max_ind and max_val
              output_shape[1] = 2
  
      self.add_output_shape(op, [output_shape])
  
imports: import math
  import numpy
  import six
  from transform.transformer import Transformer
  from transform.base_converter import DataFormat
  from transform.base_converter import MaceOp
  from transform.base_converter import MaceKeyword
  from transform.base_converter import ConverterUtil
  from utils.util import mace_check
label_desc: Use image node_id 19 for calling the ShapeInference obj's underlying member method code with example usage: obj.infer_shape_argmax(op) without return types