def scalar_op_helper(node, op_name, **kwargs):
    """Helper function for scalar arithmetic operations"""
    name, input_nodes, attrs = get_inputs(node, kwargs)
    from onnx import numpy_helper
    input_type = kwargs["in_type"]
    scalar_value = np.array([attrs.get("scalar", 1)],
                            dtype=onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[input_type])

    initializer = kwargs["initializer"]
    flag = True
    # If the input value is in initializer, just multiply with scalar input
    # and create a new initializer
    for i in initializer:
        if i.name == input_nodes[0]:
            if op_name == 'Mul':
                new_initializer = numpy_helper.to_array(i) * scalar_value[0]
            elif op_name == 'Sub':
                if name.startswith("_rminusscalar"):
                    new_initializer = scalar_value[0] - numpy_helper.to_array(i)
                else:
                    new_initializer = numpy_helper.to_array(i) - scalar_value[0]
            elif op_name == 'Add':
                new_initializer = numpy_helper.to_array(i) + scalar_value[0]
            elif op_name == 'Div':
                if name.startswith("_rdivscalar"):
                    new_initializer = scalar_value[0] / numpy_helper.to_array(i)
                else:
                    new_initializer = numpy_helper.to_array(i) / scalar_value[0]
            elif op_name == 'Pow':
                new_initializer = numpy_helper.to_array(i) ** scalar_value[0]
            flag = False
            break

    # else create a new tensor of the scalar value, add it in initializer
    if flag is True:
        dims = np.shape(scalar_value)

        scalar_op_name = "scalar_op" + str(kwargs["idx"])
        tensor_node = onnx.helper.make_tensor_value_info(scalar_op_name, input_type, dims)

        initializer.append(
            onnx.helper.make_tensor(
                name=scalar_op_name,
                data_type=input_type,
                dims=dims,
                vals=scalar_value,
                raw=False,
            )
        )

        mul_node = onnx.helper.make_node(
            op_name,
            [input_nodes[0], scalar_op_name],
            [name],
            name=name
        )

        return [tensor_node, mul_node]
    else:
        data_type = onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[new_initializer.dtype]
        dims = np.shape(new_initializer)

        new_a_node = input_nodes[0] + str(kwargs["idx"])
        tensor_node = onnx.helper.make_tensor_value_info(new_a_node, data_type, dims)

        initializer.append(
            onnx.helper.make_tensor(
                name=new_a_node,
                data_type=data_type,
                dims=dims,
                vals=new_initializer,
                raw=False,
            )
        )
        return [tensor_node]
def scalar_op_helper ( node , op_name , ** kwargs ) : name , input_nodes , attrs = get_inputs ( node , kwargs ) from onnx import numpy_helper input_type = kwargs [ "in_type" ] scalar_value = np . array ( [ attrs . get ( "scalar" , 1 ) ] , dtype = onnx . mapping . TENSOR_TYPE_TO_NP_TYPE [ input_type ] ) initializer = kwargs [ "initializer" ] flag = True for i in initializer : if i . name == input_nodes [ 0 ] : if op_name == 'Mul' : new_initializer = numpy_helper . to_array ( i ) * scalar_value [ 0 ] elif op_name == 'Sub' : if name . startswith ( "_rminusscalar" ) : new_initializer = scalar_value [ 0 ] - numpy_helper . to_array ( i ) else : new_initializer = numpy_helper . to_array ( i ) - scalar_value [ 0 ] elif op_name == 'Add' : new_initializer = numpy_helper . to_array ( i ) + scalar_value [ 0 ] elif op_name == 'Div' : if name . startswith ( "_rdivscalar" ) : new_initializer = scalar_value [ 0 ] / numpy_helper . to_array ( i ) else : new_initializer = numpy_helper . to_array ( i ) / scalar_value [ 0 ] elif op_name == 'Pow' : new_initializer = numpy_helper . to_array ( i ) ** scalar_value [ 0 ] flag = False break if flag is True : dims = np . shape ( scalar_value ) scalar_op_name = "scalar_op" + str ( kwargs [ "idx" ] ) tensor_node = onnx . helper . make_tensor_value_info ( scalar_op_name , input_type , dims ) initializer . append ( onnx . helper . make_tensor ( name = scalar_op_name , data_type = input_type , dims = dims , vals = scalar_value , raw = False , ) ) mul_node = onnx . helper . make_node ( op_name , [ input_nodes [ 0 ] , scalar_op_name ] , [ name ] , name = name ) return [ tensor_node , mul_node ] else : data_type = onnx . mapping . NP_TYPE_TO_TENSOR_TYPE [ new_initializer . dtype ] dims = np . shape ( new_initializer ) new_a_node = input_nodes [ 0 ] + str ( kwargs [ "idx" ] ) tensor_node = onnx . helper . make_tensor_value_info ( new_a_node , data_type , dims ) initializer . append ( onnx . helper . make_tensor ( name = new_a_node , data_type = data_type , dims = dims , vals = new_initializer , raw = False , ) ) return [ tensor_node ]
Helper function for scalar arithmetic operations