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

Handle Scalars in TernaryOps and Where. (#3509)

Browse source 
Handle Scalars in TernaryOps and Where.
This commit is contained in:
Dmitri Smirnov 2020-04-13 16:24:35 -07:00 committed by GitHub
parent cbe30f3e19
commit efd9b92482
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 36 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
onnxruntime/core/providers/cpu/tensor/utils.h
View file

@ -25,6 +25,11 @@ struct TensorPitches : std::vector<int64_t> {
if (gsl::narrow_cast<ptrdiff_t>(padded_rank) < 0)
return false;
// Guard against Scalars
if (pitch_rank == 0) {
return true;
}
*(p.rbegin()) = 1; // The innermost axis is 1 (single values)
if (tensor_rank > 1) {
for (size_t i = tensor_rank - 1; i-- > 0;) {

42
onnxruntime/core/providers/cuda/tensor/where.cc
View file

@ -94,40 +94,30 @@ struct TernaryElementwisePreparation {
output_rank_or_simple_broadcast = out_rank;
if (a_shape != output_shape) {
TensorPitches a_pitches(a_shape.GetDims());
a_padded_strides.size_ = out_rank;
auto offset = out_rank - a_rank;
for (auto i = offset; i < out_rank; ++i) {
// the stride for broadcast dimension is kept as 0
if (a_shape.GetDims()[i - offset] != 1) {
a_padded_strides[i] = a_pitches[i];
auto padder = [out_rank](int32_t rank, const TensorShape& shape, TArray<int64_t>& padded_strides) {
padded_strides.size_ = out_rank;
if (rank > 0) {
TensorPitches pitches(shape.GetDims());
auto offset = out_rank - rank;
for (auto i = offset; i < out_rank; ++i) {
// the stride for broadcast dimension is kept as 0
if (shape.GetDims()[i - offset] != 1) {
padded_strides[i] = pitches[i - offset];
}
}
}
};
if (a_shape != output_shape) {
padder(a_rank, a_shape, a_padded_strides);
}
if (b_shape != output_shape) {
TensorPitches b_pitches(b_shape.GetDims());
b_padded_strides.size_ = out_rank;
auto offset = out_rank - b_rank;
for (auto i = offset; i < out_rank; ++i) {
// the stride for broadcast dimension is kept as 0
if (b_shape.GetDims()[i - offset] != 1) {
b_padded_strides[i] = b_pitches[i];
}
}
padder(b_rank, b_shape, b_padded_strides);
}
if (c_shape != output_shape) {
TensorPitches c_pitches(c_shape.GetDims());
c_padded_strides.size_ = out_rank;
auto offset = out_rank - c_rank;
for (auto i = offset; i < out_rank; ++i) {
// the stride for broadcast dimension is kept as 0
if (c_shape.GetDims()[i - offset] != 1) {
c_padded_strides[i] = c_pitches[i];
}
}
padder(c_rank, c_shape, c_padded_strides);
}
TensorPitches output_pitches(output_shape.GetDims());

15
onnxruntime/test/providers/cpu/tensor/where_op_test.cc
View file

@ -122,5 +122,20 @@ TEST(WhereOpTest, BroadcastDimWithZero) {
// exclude NGraph as this isn't handled by that EP
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider});
}
TEST(WhereOpTest, BroadcastWithScalar) {
OpTester test{kOpName, kOpVersion};
test.AddInput<bool>("condition", {3}, {true, false, true});
test.AddInput<int64_t>("X", {1, 3}, {1, 2, 3});
test.AddInput<int64_t>("Y", {}, {1});
test.AddOutput<int64_t>("output", {1, 3}, {1, 1, 3});
// exclude NGraph as this isn't handled by that EP
//test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider});
test.Run();
}
} // namespace test
} // namespace onnxruntime