code: def hardmax(attrs, inputs, proto_obj):
      
      input_tensor_data = proto_obj.model_metadata.get('input_tensor_data')[0]
      input_shape = input_tensor_data[1]
  
      axis = int(attrs.get('axis', 1))
      axis = axis if axis >= 0 else len(input_shape) + axis
  
      if axis == len(input_shape) - 1:
          amax = symbol.argmax(inputs[0], axis=-1)
          one_hot = symbol.one_hot(amax, depth=input_shape[-1])
          return one_hot, attrs, inputs
  
      # since reshape doesn't take a tensor for shape,
      # computing with np.prod. This needs to be changed to
      # to use mx.sym.prod() when mx.sym.reshape() is fixed.
      # (https://github.com/apache/incubator-mxnet/issues/10789)
      new_shape = (int(np.prod(input_shape[:axis])),
                   int(np.prod(input_shape[axis:])))
      reshape_op = symbol.reshape(inputs[0], new_shape)
      amax = symbol.argmax(reshape_op, axis=-1)
      one_hot = symbol.one_hot(amax, depth=new_shape[-1])
      hardmax_op = symbol.reshape(one_hot, input_shape)
      return hardmax_op, attrs, inputs
docs: Returns batched one-hot vectors.