Create CastLike tests

microsoft · Jan 30, 2024 · d603d61 · d603d61
1 parent a4d0848
commit d603d61
Showing 1 changed file with 63 additions and 0 deletions.
diff --git a/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc b/onnxruntime/test/providers/cpu/tensor/cast_op_test.cc
@@ -12,6 +12,7 @@
 #include "core/framework/data_types_internal.h"
 
 #include "test/common/cuda_op_test_utils.h"
+#include "test/common/tensor_op_test_utils.h"
 #include "test/providers/provider_test_utils.h"
 
 namespace onnxruntime {
@@ -276,5 +277,67 @@ TEST(CastOpTest, ToFloat8E5M2FNUZ) {
 
 #endif
 
+TEST(CastLikeOpTest, CastFromFloatToInt32) {
+  OpTester test("CastLike", 18);
+  std::vector<int64_t> shape = {2, 3};
+  std::vector<float> input = {1.0f, 2.0f, 3.0f,
+                              4.0f, 5.0f, 6.0f};
+  std::vector<int32_t> output = {1, 2, 3,
+                                 4, 5, 6};
+
+  test.AddInput<float>("input", shape, input.data(), input.size());
+  test.AddInput<int32_t>("like", shape, output.data(), output.size());
+  test.AddOutput<int32_t>("output", shape, output.data(), output.size());
+
+  std::unordered_set<std::string> excluded_provider_types{kTensorrtExecutionProvider};
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider});
+}
+
+TEST(CastLikeOpTest, CastFromInt32ToFloat) {
+  OpTester test("CastLike", 18);
+  std::vector<int64_t> shape = {2, 3};
+  std::vector<int32_t> input = {1, 2, 3,
+                                4, 5, 6};
+  std::vector<float> output = {1.0f, 2.0f, 3.0f,
+                               4.0f, 5.0f, 6.0f};
+
+  test.AddInput<int32_t>("input", shape, input.data(), input.size());
+  test.AddInput<float>("like", shape, output.data(), output.size());
+  test.AddOutput<float>("output", shape, output.data(), output.size());
+
+  std::unordered_set<std::string> excluded_provider_types{kTensorrtExecutionProvider, kOpenVINOExecutionProvider};
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", excluded_provider_types);
+}
+
+TEST(CastLikeOpTest, CastFromFloatToMLFloat16) {
+  OpTester test("CastLike", 18);
+  std::vector<int64_t> shape = {2, 3};
+  std::vector<float> input = {1.5f, 2.5f, 3.5f,
+                              4.5f, 5.5f, 6.5f};
+  std::vector<MLFloat16> output = ToFloat16(input);
+
+  test.AddInput<float>("input", shape, input.data(), input.size());
+  test.AddInput<MLFloat16>("like", shape, output.data(), output.size());
+  test.AddOutput<MLFloat16>("output", shape, output.data(), output.size());
+
+  std::unordered_set<std::string> excluded_provider_types{kTensorrtExecutionProvider, kOpenVINOExecutionProvider};
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", excluded_provider_types);
+}
+
+TEST(CastLikeOpTest, CastFromMLFloat16ToFloat) {
+  OpTester test("CastLike", 18);
+  std::vector<int64_t> shape = {2, 3};
+  std::vector<float> output = {1.5f, 2.5f, 3.5f,
+                               4.5f, 5.5f, 6.5f};
+  std::vector<MLFloat16> input = ToFloat16(output);
+
+  test.AddInput<MLFloat16>("input", shape, input.data(), input.size());
+  test.AddInput<float>("like", shape, output.data(), output.size());
+  test.AddOutput<float>("output", shape, output.data(), output.size());
+
+  std::unordered_set<std::string> excluded_provider_types{kTensorrtExecutionProvider, kOpenVINOExecutionProvider};
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", excluded_provider_types);
+}
+
 }  // namespace test
 }  // namespace onnxruntime