Function _replace_constant_of_shape_value

onnx/tools/replace_constants.py:157–207

Replaces all fill value of all nodes *ConstantOfShape*.

(
    onx: GraphProto | FunctionProto, value_constant_of_shape: float
)

def _replace_constant_of_shape_value(
    onx: GraphProto | FunctionProto, value_constant_of_shape: float
) -> GraphProto | FunctionProto:
    """Replaces all fill value of all nodes *ConstantOfShape*."""
    if isinstance(onx, GraphProto):
        nodes = list(onx.node)
    elif isinstance(onx, FunctionProto):
        nodes = list(onx.node)
    else:
        raise TypeError(f"Not implemented for type {type(onx)}.")

    existing_names = set()
    for node in nodes:
        existing_names |= set(node.input)
        existing_names |= set(node.output)

    update = {}
    for inode, node in enumerate(nodes):
        if node.op_type != "ConstantOfShape":
            continue
        tensor = node.attribute[0].t
        new_tensor = make_tensor(
            tensor.name, tensor.data_type, [1], [value_constant_of_shape]
        )
        new_node = make_node("ConstantOfShape", node.input, node.output)
        att = make_attribute(node.attribute[0].name, value=new_tensor)
        new_node.attribute.append(att)
        update[inode] = new_node

    for inode, up in update.items():
        nodes[inode] = up

    if isinstance(onx, GraphProto):
        return make_graph(
            nodes,
            onx.name,
            onx.input,
            onx.output,
            initializer=onx.initializer,
            sparse_initializer=onx.sparse_initializer,
        )
    if isinstance(onx, FunctionProto):
        return make_function(
            onx.domain,
            onx.name,
            onx.input,
            onx.output,
            nodes,
            opset_imports=onx.opset_import,
        )
    raise TypeError(f"Not implemented for type {type(onx)}.")


def replace_initializer_by_constant_of_shape(  # noqa: PLR0911