diff --git a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc index 970de353f5..6be9ac8a60 100644 --- a/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc +++ b/onnxruntime/core/optimizer/embed_layer_norm_fusion.cc @@ -65,6 +65,193 @@ static bool CheckInput(NodeArg* input, const logging::Logger& logger) { return true; } +static bool MatchPositionEmbeddingSubgraph1( + Graph& graph, + Node& position_gather_node, + NodeArg* input_ids, + const logging::Logger& logger, + std::vector& matched_edges) { + // Match two paths. + // Match Shape --> Expand path if needed. + std::vector position_parent_nodes; + std::vector position_embedding_path_symbolic{ + {0, 1, "Expand", {8}, kOnnxDomain}, + {0, 1, "Shape", {1}, kOnnxDomain}}; + std::vector edges; + if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { + return false; + } + if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } + + // Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand + Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); + Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); + std::vector pg_parent_path{ + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Cast", {9}, kOnnxDomain}, + {0, 0, "Squeeze", {1}, kOnnxDomain}, + {0, 0, "Transpose", {1}, kOnnxDomain}, + {0, 0, "NonZero", {9}, kOnnxDomain}, + {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}, + }; + matched_edges = edges; + + if (!graph_utils::FindPath(expand_node, true, pg_parent_path, edges, logger)) { + return false; + } + for (size_t i = 0; i < edges.size(); i++) { + if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + return false; + } + } + // Check if the second input of the Gather node in the path has a constant input of 1 + Node& gather_node = *graph.GetNode(edges[edges.size() - 2]->GetNode().Index()); + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) { + DEBUG_LOG("Second input of Gather should be a constant with value 1. "); + return false; + } + + // Check if the parent of "shape" is the input_ids + Node& shape_node_2 = *graph.GetNode(edges[edges.size() - 1]->GetNode().Index()); + if (shape_node_1.MutableInputDefs()[0] != input_ids || + shape_node_2.MutableInputDefs()[0] != input_ids) { + return false; + } + + matched_edges.insert(matched_edges.end(), edges.begin(), edges.end()); + return true; +} + +/** Match subgraph like the following: + (input_ids) + / \ + Shape Shape + | | + Gather (indice=0) Gather (indice=1)--+ + | | | + Unsqueeze Unsqueeze Cast + \ / | + \ / Range(start=0, delta=1) + \ / | + Concat Unsqueeze + | | + +--|----------------------------+ + | | + Expand + | + Gather + + Note that position gather node is the node in the bottom of above sub-graph. +*/ + +static bool MatchPositionEmbeddingSubgraph2( + Graph& graph, + Node& position_gather_node, + NodeArg* input_ids, + const logging::Logger& logger, + std::vector& matched_edges) { + + // Match Gather <-- Expand <-- Unsqueeze <-- Range <-- Cast <-- Gather <-- Shape + std::vector position_parent_nodes; + std::vector position_embedding_path_symbolic{ + {0, 1, "Expand", {8}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Range", {11}, kOnnxDomain}, + {0, 1, "Cast", {9}, kOnnxDomain}, + {0, 0, "Gather", {11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}, + }; + std::vector edges; + if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { + DEBUG_LOG("Failed to find path 1."); + return false; + } + for (size_t i = 0; i < edges.size(); i++) { + if (edges[i]->GetNode().GetOutputEdgesCount() != (i == 4 ? 2 : 1)) { + DEBUG_LOG("Output edge count not expected for nodes in path 1."); + return false; + } + } + matched_edges = edges; + + Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); + Node& range_node = *graph.GetNode(edges[2]->GetNode().Index()); + Node& gather_node_1 = *graph.GetNode(edges[4]->GetNode().Index()); + Node& shape_node_1 = *graph.GetNode(edges[5]->GetNode().Index()); + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(range_node.InputDefs()[0]), int64_t(0), true)) { + DEBUG_LOG("The first input of Range should be a constant with value 0."); + return false; + } + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(range_node.InputDefs()[2]), int64_t(1), true)) { + DEBUG_LOG("The third input of Range should be a constant with value 1."); + return false; + } + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node_1.InputDefs()[1]), int64_t(1), true)) { + DEBUG_LOG("The second input of Gather in path1 should be a constant with value 1."); + return false; + } + + std::vector expand_parent_path1{ + {0, 1, "Concat", {11}, kOnnxDomain}, + {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain}, + {0, 0, "Shape", {1}, kOnnxDomain}, + }; + if (!graph_utils::FindPath(expand_node, true, expand_parent_path1, edges, logger)) { + DEBUG_LOG("Failed to find path 2."); + return false; + } + for (size_t i = 0; i < edges.size(); i++) { + if (edges[i]->GetNode().GetOutputEdgesCount() != 1) { + DEBUG_LOG("Output edge count not expected for nodes in path 2."); + return false; + } + } + + Node& concat_node = *graph.GetNode(edges[0]->GetNode().Index()); + Node& gather_node_0 = *graph.GetNode(edges[2]->GetNode().Index()); + Node& shape_node_0 = *graph.GetNode(edges[3]->GetNode().Index()); + if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node_0.InputDefs()[1]), int64_t(0), true)) { + DEBUG_LOG("Second input of Gather in path2 should be a constant with value 0."); + return false; + } + matched_edges.insert(matched_edges.end(), edges.begin(), edges.end()); + + std::vector concat_parent_path{ + {0, 1, "Unsqueeze", {1, 11}, kOnnxDomain}, + {0, 0, "Gather", {1, 11}, kOnnxDomain} + }; + if (!graph_utils::FindPath(concat_node, true, concat_parent_path, edges, logger)) { + DEBUG_LOG("Failed to find path 3."); + return false; + } + // Two paths share the gather node (with second input indices==1) + if (edges[1]->GetNode().Index() != gather_node_1.Index()) { + DEBUG_LOG(" Gather nodes in path 1 and 3 expected to be same node."); + return false; + } + if (edges[0]->GetNode().GetOutputEdgesCount() != 1) { + DEBUG_LOG("Output edge count not expected for nodes in path 3."); + return false; + } + + // Check if the two paths of position gather lead to the same input. + if (shape_node_0.MutableInputDefs()[0] != input_ids || + shape_node_1.MutableInputDefs()[0] != input_ids) { + DEBUG_LOG("The parent of two shape nodes are expected to be input_ids."); + return false; + } + + // Do not add the gather node since it has been added in another path. + matched_edges.push_back(edges[0]); + return true; +} + /** Embed Layer Normalization will fuse embeddings and mask processing into one node : The embeddings before conversion: @@ -176,22 +363,19 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l continue; } + NodeArg* input_ids = word_gather_node.MutableInputDefs()[1]; + // Check the second input of position gather. If it's not initializer, check for two paths. - Node* p_expand_node = nullptr; - Node* p_shape_node = nullptr; std::vector pg_edges; - bool isValidEmbedSubNode = true; if (graph_utils::IsConstantInitializer(graph, position_gather_node.MutableInputDefs()[1]->Name())) { - // Check if the second input of position gather is a tensor with values evenly spaced by 1 starting from 0. + // Check if the second input of position gather is a tensor with values evenly spaced by 1 starting from 0. std::vector data; auto expected_shape = word_gather_node.MutableInputDefs()[1]->Shape(); - if (!optimizer_utils::AppendTensorFromInitializer(graph, *(position_gather_node.MutableInputDefs()[1]), data) - || !utils::HasDimValue(expected_shape->dim()[0]) - || !utils::HasDimValue(expected_shape->dim()[1]) - || static_cast(data.size()) != expected_shape->dim()[0].dim_value() * expected_shape->dim()[1].dim_value()) { + if (!optimizer_utils::AppendTensorFromInitializer(graph, *(position_gather_node.MutableInputDefs()[1]), data) || !utils::HasDimValue(expected_shape->dim()[0]) || !utils::HasDimValue(expected_shape->dim()[1]) || static_cast(data.size()) != expected_shape->dim()[0].dim_value() * expected_shape->dim()[1].dim_value()) { continue; } int64_t expected_value = 0; + bool isValidEmbedSubNode = true; for (size_t i = 0; i < data.size(); i++) { if (data[i] != expected_value) { isValidEmbedSubNode = false; @@ -202,68 +386,19 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l expected_value = 0; } } + if (!isValidEmbedSubNode) { + continue; + } } else { - // Match two paths. - // Match Shape --> Expand path if needed. - std::vector position_parent_nodes; - std::vector position_embedding_path_symbolic{ - {0, 1, "Expand", {8}, kOnnxDomain}, - {0, 1, "Shape", {1}, kOnnxDomain}}; - if (!graph_utils::FindPath(position_gather_node, true, position_embedding_path_symbolic, edges, logger)) { - continue; - } - if (edges[0]->GetNode().GetOutputEdgesCount() != 1 && edges[1]->GetNode().GetOutputEdgesCount() != 1) { - continue; - } - p_expand_node = graph.GetNode(edges[0]->GetNode().Index()); - p_shape_node = graph.GetNode(edges[1]->GetNode().Index()); - // Match Shape --> Gather --> Unsqueeze --> ConstantOfShape --> NonZero --> Transpose --> Squeeze --> Cast --> Unsqueeze --> Expand - Node& expand_node = *graph.GetNode(edges[0]->GetNode().Index()); - Node& shape_node_1 = *graph.GetNode(edges[1]->GetNode().Index()); - std::vector pg_parent_path{ - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Cast", {9}, kOnnxDomain}, - {0, 0, "Squeeze", {1}, kOnnxDomain}, - {0, 0, "Transpose", {1}, kOnnxDomain}, - {0, 0, "NonZero", {9}, kOnnxDomain}, - {0, 0, "ConstantOfShape", {9}, kOnnxDomain}, - {0, 0, "Unsqueeze", {1, 11}, kOnnxDomain}, - {0, 0, "Gather", {1, 11}, kOnnxDomain}, - {0, 0, "Shape", {1}, kOnnxDomain}, - }; - if (!graph_utils::FindPath(expand_node, true, pg_parent_path, pg_edges, logger)) { - continue; - } - for (size_t i = 0; i < pg_edges.size(); i++) { - if (pg_edges[i]->GetNode().GetOutputEdgesCount() != 1) { - isValidEmbedSubNode = false; - break; + if (!MatchPositionEmbeddingSubgraph1(graph, position_gather_node, input_ids, logger, pg_edges)) { + pg_edges.clear(); + if (!MatchPositionEmbeddingSubgraph2(graph, position_gather_node, input_ids, logger, pg_edges)) { + continue; } } - // Check if the second input of the Gather node in the path has a constant input of 1 - Node& gather_node = *graph.GetNode(pg_edges[pg_edges.size() - 2]->GetNode().Index()); - if (!optimizer_utils::IsInitializerWithExpectedValue(graph, *(gather_node.InputDefs()[1]), int64_t(1), true)) { - DEBUG_LOG("Second input of Gather should be a constant with value 1. "); - - continue; - } - // Check if the two paths of position gather lead to the same input. - Node& shape_node_2 = *graph.GetNode(pg_edges[pg_edges.size() - 1]->GetNode().Index()); - if (shape_node_1.MutableInputDefs()[0] != shape_node_2.MutableInputDefs()[0]) { - continue; - } - // Check if the parent of "shape" is the parent of "word gather" - if (shape_node_1.MutableInputDefs()[0] != word_gather_node.MutableInputDefs()[1]) { - continue; - } - - } - if (!isValidEmbedSubNode) { - continue; } // Get input "input_ids" from node. - NodeArg* input_ids = word_gather_node.MutableInputDefs()[1]; if (!CheckInput(input_ids, logger)) { DEBUG_LOG("Input id is not valid. "); continue; @@ -283,32 +418,30 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l continue; } - if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != - utils::GetTensorShapeFromTensorShapeProto(*(segment_ids->Shape()))) { + if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != + utils::GetTensorShapeFromTensorShapeProto(*(segment_ids->Shape()))) { DEBUG_LOG("Input_ids and segment id should have the same shape. "); continue; } - if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != - utils::GetTensorShapeFromTensorShapeProto(*(mask->Shape()))) { + if (utils::GetTensorShapeFromTensorShapeProto(*(input_ids->Shape())) != + utils::GetTensorShapeFromTensorShapeProto(*(mask->Shape()))) { DEBUG_LOG("Input_ids and mask should have the same shape. "); continue; } NodeArg* gamma = layer_norm_node.MutableInputDefs()[1]; NodeArg* beta = layer_norm_node.MutableInputDefs()[2]; - if (gamma->Shape() == nullptr - || gamma->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { + if (gamma->Shape() == nullptr || gamma->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { DEBUG_LOG("Gamma should be of shape (hidden_size). "); continue; } - if (beta->Shape() == nullptr - || beta->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { + if (beta->Shape() == nullptr || beta->Shape()->dim()[0].dim_value() != word_embedding->Shape()->dim()[1].dim_value()) { DEBUG_LOG("Beta should be of shape (hidden_size). "); continue; } - // Cast input_ids, segment_ids, and mask to int32 if needed. + // Cast input_ids, segment_ids, and mask to int32 if needed. input_ids = CastToInt32(graph, input_ids, layer_norm_node.GetExecutionProviderType()); segment_ids = CastToInt32(graph, segment_ids, layer_norm_node.GetExecutionProviderType()); mask = CastToInt32(graph, mask, layer_norm_node.GetExecutionProviderType()); @@ -339,10 +472,6 @@ Status EmbedLayerNormFusion::ApplyImpl(Graph& graph, bool& modified, int graph_l for (size_t i = 0; i < pg_edges.size(); i++) { nodes_to_remove.push_back(pg_edges[i]->GetNode().Index()); } - if (p_shape_node != nullptr && p_expand_node != nullptr) { - nodes_to_remove.push_back(p_shape_node->Index()); - nodes_to_remove.push_back(p_expand_node->Index()); - } nodes_to_remove.push_back(word_gather_node.Index()); nodes_to_remove.push_back(position_gather_node.Index()); nodes_to_remove.push_back(segment_gather_node.Index()); diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index bd7665fded..ad108f64df 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -1328,6 +1328,32 @@ TEST(GraphTransformationTests, EmbedLayerNormFusionFormat2) { ASSERT_TRUE(op_to_count["SkipLayerNormalization"] == 0); ASSERT_TRUE(op_to_count["EmbedLayerNormalization"] == 1); } + +TEST(GraphTransformationTests, EmbedLayerNormFusionFormat3) { + auto model_uri = MODEL_FOLDER "fusion/embed_layer_norm_format3.onnx"; + std::shared_ptr p_model; + ASSERT_TRUE(Model::Load(model_uri, p_model, nullptr, DefaultLoggingManager().DefaultLogger()).IsOK()); + Graph& graph = p_model->MainGraph(); + + onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; + graph_transformation_mgr.Register(onnxruntime::make_unique(), TransformerLevel::Level2); + auto ret = graph_transformation_mgr.ApplyTransformers(graph, TransformerLevel::Level2, DefaultLoggingManager().DefaultLogger()); + ASSERT_TRUE(ret.IsOK()); + + std::map op_to_count = CountOpsInGraph(graph); + EXPECT_EQ(op_to_count["Shape"], 0); + EXPECT_EQ(op_to_count["Expand"], 0); + EXPECT_EQ(op_to_count["Gather"], 0); + EXPECT_EQ(op_to_count["Unsqueeze"], 0); + EXPECT_EQ(op_to_count["LayerNormalization"], 0); + EXPECT_EQ(op_to_count["SkipLayerNormalization"], 0); + EXPECT_EQ(op_to_count["ReduceSum"], 0); + EXPECT_EQ(op_to_count["MatMul"], 1); + EXPECT_EQ(op_to_count["Add"], 2); + EXPECT_EQ(op_to_count["Cast"], 3); + EXPECT_EQ(op_to_count["Attention"], 1); + EXPECT_EQ(op_to_count["EmbedLayerNormalization"], 1); +} #endif } // namespace test diff --git a/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format3.onnx b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format3.onnx new file mode 100644 index 0000000000..e58ffa62a1 Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_format3.onnx differ diff --git a/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_gen.py b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_gen.py new file mode 100644 index 0000000000..b1fdc697d0 --- /dev/null +++ b/onnxruntime/test/testdata/transform/fusion/embed_layer_norm_gen.py @@ -0,0 +1,75 @@ +import onnx +from onnx import helper +from onnx import TensorProto +from enum import Enum + +class Format(Enum): + Format1=1, + Format2=2, + Format3=3 + +def GenerateModel(format, model_name): + nodes = [ # LayerNorm subgraph + helper.make_node("Shape", ["input_ids"], ["shape1_out"], "shape1"), + helper.make_node("Gather", ["shape1_out", "indices_0"], ["gather0_out"], "gather0"), + helper.make_node("Unsqueeze", ["gather0_out"], ["unsqueeze0_out"], "unsqueeze0", axes=[0]), + helper.make_node("Shape", ["input_ids"], ["shape2_out"], "shape2"), + helper.make_node("Gather", ["shape2_out", "indices_1"], ["gather1_out"], "gather1"), + helper.make_node("Unsqueeze", ["gather1_out"], ["unsqueeze1_out"], "unsqueeze1", axes=[0]), + helper.make_node("Concat", ["unsqueeze0_out", "unsqueeze1_out"], ["concat_out"], "concat", axis=0), + helper.make_node("Cast", ["gather1_out"], ["cast_out"], "cast", to=7), + helper.make_node("Range", ["start_0", "cast_out", "delta_1"], ["range_out"], "range"), + helper.make_node("Unsqueeze", ["range_out"], ["unsqueeze2_out"], "unsqueeze2", axes=[0]), + helper.make_node("Expand", ["unsqueeze2_out", "concat_out"], ["expand_out"], "expand"), + helper.make_node("Gather", ["pos_embed", "expand_out"], ["pos_gather_out"], "pos_gather"), + helper.make_node("Gather", ["word_embed", "input_ids"], ["word_gather_out"], "word_gather"), + helper.make_node("Add", ["word_gather_out", "pos_gather_out"], ["word_add_pos_out"], "word_add_pos"), + helper.make_node("Gather", ["seg_embed", "segment_ids"], ["seg_gather_out"], "seg_gather"), + helper.make_node("Add", ["word_add_pos_out", "seg_gather_out"], ["add3_out"], "add3"), + helper.make_node("LayerNormalization", ["add3_out", "layer_norm_weight", "layer_norm_bias"], ["layernorm_out"], "layernorm", axis=-1, epsion=0.000009999999747378752), + helper.make_node("Cast", ["input_mask"], ["mask_cast_out"], "mask_cast", to=6), + helper.make_node("ReduceSum", ["mask_cast_out"], ["mask_index_out"], "mask_index", axes=[1], keepdims=0), + helper.make_node("Attention", ["layernorm_out", "qkv_weights", "qkv_bias", "mask_index_out"], ["att_out"], "att", domain="com.microsoft", num_heads=2), + helper.make_node("MatMul", ["att_out", "matmul_weight"], ["matmul_out"], "matmul"), + helper.make_node("Add", ["matmul_out", "add_bias"], ["add_out"], "add"), + helper.make_node("Add", ["add_out", "layernorm_out"], ["add2_out"], "add2") + + ] + + # hidden_size=4, num_heads=2, max_seq_length=3 + initializers = [ # initializers + helper.make_tensor('indices_0', TensorProto.INT64, [], [0]), + helper.make_tensor('indices_1', TensorProto.INT64, [], [1]), + helper.make_tensor('start_0', TensorProto.INT64, [], [0]), + helper.make_tensor('delta_1', TensorProto.INT64, [], [1]), + helper.make_tensor('word_embed', TensorProto.FLOAT, [2, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]), + helper.make_tensor('pos_embed', TensorProto.FLOAT, [4, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]), + helper.make_tensor('seg_embed', TensorProto.FLOAT, [2, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]), + helper.make_tensor('layer_norm_weight', TensorProto.FLOAT, [4], [1.0, 2.0, 3.0, 4.0]), + helper.make_tensor('layer_norm_bias', TensorProto.FLOAT, [4], [0.1, 0.2, 0.3, 0.4]), + + helper.make_tensor('qkv_weights', TensorProto.FLOAT, [4, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]), + helper.make_tensor('qkv_bias', TensorProto.FLOAT, [4], [0.1, 0.2, 0.3, 0.4]), + + helper.make_tensor('matmul_weight', TensorProto.FLOAT, [4, 4], [1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0, 1.0, 2.0, 3.0, 4.0]), + helper.make_tensor('add_bias', TensorProto.FLOAT, [4], [0.1, 0.2, 0.3, 0.4]), + ] + + graph = helper.make_graph( + nodes, + "EmbedLayerNorm_format3", #name + [ # inputs + helper.make_tensor_value_info('input_ids', TensorProto.INT64, ['batch', 3]), + helper.make_tensor_value_info('segment_ids', TensorProto.INT64, ['batch', 3]), + helper.make_tensor_value_info('input_mask', TensorProto.INT64, ['batch', 3]), + ], + [ # outputs + helper.make_tensor_value_info('add2_out', TensorProto.FLOAT, ['batch', 3, 4]), + ], + initializers + ) + + model = helper.make_model(graph) + onnx.save(model, model_name) + +GenerateModel(Format.Format3, 'embed_layer_norm_format3.onnx') \ No newline at end of file