mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
Fix regression introduced in cast transformer (#4658)
This commit is contained in:
parent
f90a2d46ae
commit
382f94c95c
3 changed files with 27 additions and 3 deletions
|
|
@ -254,7 +254,10 @@ Status InsertCastTransformer::ApplyImpl(onnxruntime::Graph& graph, bool& modifie
|
|||
ORT_ENFORCE(dtype_attribute->second.has_i(),
|
||||
"InsertCastTransformer works on the assumption that `dtype` attribute holds an integer.");
|
||||
|
||||
dtype_attribute->second.set_i(TensorProto_DataType_FLOAT);
|
||||
// Modify the dtype attribute (which defines the output type) to FLOAT if it is FLOAT16.
|
||||
if (dtype_attribute->second.i() == TensorProto_DataType_FLOAT16) {
|
||||
dtype_attribute->second.set_i(TensorProto_DataType_FLOAT);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -134,8 +134,9 @@ TEST(TransformerTest, ThreeInARowRemoval) {
|
|||
ASSERT_TRUE(op_to_count["Cast"] == 2);
|
||||
}
|
||||
|
||||
// test a case where the ONNX inferred type (float16) is different from the type bound
|
||||
// to the output NodeArg of the "RandomNormalLike" node because of the InsertCaseTransformer
|
||||
// test a case where the ONNX inferred output type (float16) is different from the type bound
|
||||
// to the output NodeArg of the "RandomNormalLike" node (input is float16) because of the InsertCaseTransformer
|
||||
// Here the ONNX inferred output type (float16) must be made float because that is what the kernel produces
|
||||
TEST(TransformerTest, RandomNormalLikeWithFloat16Inputs) {
|
||||
auto model_uri = MODEL_FOLDER ORT_TSTR("random_normal_like_float16.onnx");
|
||||
std::shared_ptr<Model> model;
|
||||
|
|
@ -153,5 +154,25 @@ TEST(TransformerTest, RandomNormalLikeWithFloat16Inputs) {
|
|||
EXPECT_TRUE(status.IsOK()) << status;
|
||||
}
|
||||
|
||||
// A case where the ONNX inferred output type is int32 to a node that consumes float16 input
|
||||
// Here the InsertCastTransformer must not change the ONNX inferred output type and keep it
|
||||
// as is (int32)
|
||||
TEST(TransformerTest, MultinomialWithFloat16Input) {
|
||||
auto model_uri = MODEL_FOLDER ORT_TSTR("multinomial_float16.onnx");
|
||||
std::shared_ptr<Model> model;
|
||||
auto status = Model::Load(model_uri, model, nullptr, DefaultLoggingManager().DefaultLogger());
|
||||
ASSERT_TRUE(status.IsOK()) << status;
|
||||
|
||||
Graph& graph = model->MainGraph();
|
||||
InsertCastTransformer transformer("Test");
|
||||
|
||||
bool modified = false;
|
||||
status = transformer.Apply(graph, modified, DefaultLoggingManager().DefaultLogger());
|
||||
EXPECT_TRUE(status.IsOK()) << status;
|
||||
EXPECT_TRUE(modified) << "Transformer should have added some Cast nodes";
|
||||
status = graph.Resolve();
|
||||
EXPECT_TRUE(status.IsOK()) << status;
|
||||
}
|
||||
|
||||
} // namespace test
|
||||
} // namespace onnxruntime
|
||||
|
|
|
|||
BIN
onnxruntime/test/testdata/transform/multinomial_float16.onnx
vendored
Normal file
BIN
onnxruntime/test/testdata/transform/multinomial_float16.onnx
vendored
Normal file
Binary file not shown.
Loading…
Reference in a new issue