saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-27 20:02:15 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Address scalar expand to scalar case (#5854)

Browse source 
* address scalar expand to scalar case

* remove tensorrt unsupported type

* exclude tensorrt from scalar to scalar case

Co-authored-by: Randy Shuai <rashuai@microsoft.com>
This commit is contained in:
RandySheriffH 2020-11-19 11:25:53 -08:00 committed by GitHub
parent e485805b8d
commit 4bb41d610b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 26 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
onnxruntime/core/providers/cpu/tensor/expand.cc
View file

@ -72,6 +72,12 @@ Status Expand<T>::Compute(OpKernelContext* context) const {
auto* output_dims = output_shape.data();
auto output_dims_size = static_cast<int64_t>(output_shape.size());
auto max_dims_size = std::max(input_dims_size, output_dims_size);
if (0 == max_dims_size) {
*output_data = *input_data;
return Status::OK();
}
std::unique_ptr<int64_t[]> input_dim_group{new int64_t[max_dims_size]};
std::unique_ptr<int64_t[]> output_dim_group{new int64_t[max_dims_size]};
std::unique_ptr<int64_t[]> expand_dim_size{new int64_t[max_dims_size]};

20
onnxruntime/test/providers/cpu/tensor/expand_test.cc
View file

@ -163,5 +163,25 @@ TEST(ExpandOpTest, Expand_2x2x1x2x1_float) {
test.Run();
}
#ifndef USE_TENSORRT
TEST(ExpandOpTest, Expand_scalar_float) {
OpTester test("Expand", 8);
test.AddInput<float>("data_0", {}, {3.0f});
test.AddInput<int64_t>("data_1", {0}, {});
test.AddOutput<float>("result", {}, {3.0f});
test.Run();
}
#endif
TEST(ExpandOpTest, Expand_scalar_int32) {
OpTester test("Expand", 8);
test.AddInput<int32_t>("data_0", {}, {9});
test.AddInput<int64_t>("data_1", {3}, {2, 3, 4});
test.AddOutput<int32_t>("result", {2, 3, 4},
{9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9,
9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9, 9});
test.Run();
}
} //namespace test
} //namespace onnxruntime