(builder, node, graph, err)
Conversion
Conversion
Conversion
def _convert_argmax(builder, node, graph, err):  # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None

    def _add_argmax_or_argmin(input_names, output_names, **kwargs):
        input_name = input_names[0]
        output_name = output_names[0]
        if kwargs['node'].op_type == 'ArgMin':
            kwargs['builder'].add_elementwise(name=kwargs['node'].name + '_multiply_minus_1',  # type: ignore
                                    input_names=[input_name],
                                    output_name=input_name + '_multiply_minus_1',
                                    mode='MULTIPLY',
                                    alpha=-1)
            input_name += '_multiply_minus_1'
        kwargs['builder'].add_reduce(name=kwargs['node'].name,
                           input_name=input_name, output_name=output_name,
                           axis=kwargs['coreml_axis'],
                           mode='argmax')
    '''
    Conversion
    '''
    axis = node.attrs.get('axis', 0)
    keepdims = node.attrs.get('keepdims', 1)

    input_name = node.inputs[0]
    output_name = node.outputs[0]

    if _is_input_shape_mapping_defined(node, graph):
        mapp = graph.onnx_coreml_shape_mapping[node.inputs[0]]
        coreml_axis = mapp[axis]
        coreml_axis_string = "C"
        if coreml_axis == 1: # coreml_axis corresponds to the batch dimension
            return err.unsupported_op_configuration(builder, node, graph,"Cannot apply operation along Batch axis")
        if coreml_axis != 0:
            coreml_axis_string = ['C','H','W'][coreml_axis-2]
            _add_argmax_or_argmin([input_name], [output_name], builder=builder,node=node,coreml_axis=coreml_axis_string)
        else: # coreml_axis corresponds to the sequence dimension
            _add_transpose_before_after(_add_argmax_or_argmin,
                                        [input_name],
                                        [output_name],
                                        [1,0,2,3],
                                        builder=builder,node=node,coreml_axis=coreml_axis_string)

    else:
        coreml_axis_string = _get_coreml_axis([axis], builder, node, graph, err)
        if coreml_axis_string not in ['C', 'H', 'W', 'HW', 'CHW']:
            return err.unsupported_op_configuration(builder, node, graph,
                                                    "Unable to translate axes attribute to CoreML axis parameter for %s" % axis)
        _add_argmax_or_argmin([input_name], [output_name], builder=builder, node=node, coreml_axis=coreml_axis_string)

    '''
    update output shape map
    '''
    if _is_input_shape_mapping_defined(node, graph):
        mapp = graph.onnx_coreml_shape_mapping[node.inputs[0]]
        if keepdims == 1:
            graph.onnx_coreml_shape_mapping[node.outputs[0]] = mapp
        else:
            graph.onnx_coreml_shape_mapping[node.outputs[0]] = mapp[:axis] + mapp[axis+1:]
def _convert_argmax ( builder , node , graph , err ) : # type: (NeuralNetworkBuilder, Node, Graph, ErrorHandling) -> None def _add_argmax_or_argmin ( input_names , output_names , * * kwargs ) : input_name = input_names [ 0 ] output_name = output_names [ 0 ] if kwargs [ 'node' ] . op_type == 'ArgMin' : kwargs [ 'builder' ] . add_elementwise ( name = kwargs [ 'node' ] . name + '_multiply_minus_1' , # type: ignore input_names = [ input_name ] , output_name = input_name + '_multiply_minus_1' , mode = 'MULTIPLY' , alpha = - 1 ) input_name += '_multiply_minus_1' kwargs [ 'builder' ] . add_reduce ( name = kwargs [ 'node' ] . name , input_name = input_name , output_name = output_name , axis = kwargs [ 'coreml_axis' ] , mode = 'argmax' ) axis = node . attrs . get ( 'axis' , 0 ) keepdims = node . attrs . get ( 'keepdims' , 1 ) input_name = node . inputs [ 0 ] output_name = node . outputs [ 0 ] if _is_input_shape_mapping_defined ( node , graph ) : mapp = graph . onnx_coreml_shape_mapping [ node . inputs [ 0 ] ] coreml_axis = mapp [ axis ] coreml_axis_string = "C" if coreml_axis == 1 : # coreml_axis corresponds to the batch dimension return err . unsupported_op_configuration ( builder , node , graph , "Cannot apply operation along Batch axis" ) if coreml_axis != 0 : coreml_axis_string = [ 'C' , 'H' , 'W' ] [ coreml_axis - 2 ] _add_argmax_or_argmin ( [ input_name ] , [ output_name ] , builder = builder , node = node , coreml_axis = coreml_axis_string ) else : # coreml_axis corresponds to the sequence dimension _add_transpose_before_after ( _add_argmax_or_argmin , [ input_name ] , [ output_name ] , [ 1 , 0 , 2 , 3 ] , builder = builder , node = node , coreml_axis = coreml_axis_string ) else : coreml_axis_string = _get_coreml_axis ( [ axis ] , builder , node , graph , err ) if coreml_axis_string not in [ 'C' , 'H' , 'W' , 'HW' , 'CHW' ] : return err . unsupported_op_configuration ( builder , node , graph , "Unable to translate axes attribute to CoreML axis parameter for %s" % axis ) _add_argmax_or_argmin ( [ input_name ] , [ output_name ] , builder = builder , node = node , coreml_axis = coreml_axis_string ) '''
    update output shape map
    ''' if _is_input_shape_mapping_defined ( node , graph ) : mapp = graph . onnx_coreml_shape_mapping [ node . inputs [ 0 ] ] if keepdims == 1 : graph . onnx_coreml_shape_mapping [ node . outputs [ 0 ] ] = mapp else : graph . onnx_coreml_shape_mapping [ node . outputs [ 0 ] ] = mapp [ : axis ] + mapp [ axis + 1 : ]