def _assign_onnx_node_name(graph, node_names):
    """Takes in ONNX graph, and mapping from torch._C.Node to node name in exported ONNX ModelProto.

    Returns:
        graph (torch._C.Graph): A TorchScript IR Graph with ONNX nodes, where each torch._C.Node gets its name
        in exported ONNX ModelProto assigned as attribute ``onnx_name``.
    """

    def n_fn(n, b_fn, node_names):
        for b in n.blocks():
            b_fn(b, node_names)
        if n in node_names:
            n.s_("onnx_name", node_names[n])

    def b_fn(b, node_names):
        for n in b.nodes():
            n_fn(n, b_fn, node_names)

    b_fn(graph, node_names)
    return graph