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

Fix bug with Slice. Need to pass in flattened input dimensions so the initial offset into the input is calculated correctly. (#2372)

Browse source
This commit is contained in:
Scott McKay 2019-11-13 07:00:26 +10:00 committed by GitHub
parent 9e26f4de6f
commit c0d23d5ffe
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 42 additions and 14 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

38
onnxruntime/core/providers/cpu/tensor/slice.cc
View file

@ -332,23 +332,35 @@ Status SliceImpl(OpKernelContext* ctx,
auto* output = output_tensor.template MutableData<T>();
const auto* output_end = output + output_tensor.Shape().Size();
SliceIterator<T> input_iterator =
flattened_output_dims != nullptr
? SliceIterator<T>(input_tensor, *flattened_output_dims, starts, *flattened_output_dims, steps)
: SliceIterator<T>(input_tensor, starts, output_dims, steps);
if (input_iterator.SolitaryInnerStep()) {
while (output < output_end) {
output = input_iterator.CopyInnermostAxisSolitaryInnerStep(output);
auto create_output = [&output, &output_end](SliceIterator<T>& input_iterator) {
if (input_iterator.SolitaryInnerStep()) {
while (output < output_end) {
output = input_iterator.CopyInnermostAxisSolitaryInnerStep(output);
}
} else {
while (output < output_end) {
output = input_iterator.CopyInnermostAxisNonSolitaryInnerStep(output);
}
}
ORT_ENFORCE(output == output_end);
};
if (flattened_output_dims) {
// if we have flattened output dims we need to also flatten the input dims.
// as we're combining the innermost dims and keeping all values we can just copy the size of the last dim
std::vector<int64_t> flattened_input_dims(input_tensor.Shape().GetDims());
flattened_input_dims.resize(flattened_output_dims->size());
flattened_input_dims.back() = flattened_output_dims->back();
TensorShape input_shape(std::move(flattened_input_dims));
auto input_iterator = SliceIterator<T>(input_tensor, input_shape, starts, *flattened_output_dims, steps);
create_output(input_iterator);
} else {
while (output < output_end) {
output = input_iterator.CopyInnermostAxisNonSolitaryInnerStep(output);
}
auto input_iterator = SliceIterator<T>(input_tensor, starts, output_dims, steps);
create_output(input_iterator);
}
ORT_ENFORCE(output == output_end);
return Status::OK();
}

18
onnxruntime/test/providers/cpu/tensor/slice_op.test.cc
View file

@ -385,7 +385,7 @@ TEST(SliceTest, Slice3D_FlattenInnermostDimsIncopy) {
17, 6, 24},
{0, 0, 1},
{1000, 1000, 1000},
{2, 1, 0}, // reverse the axes to test that's handled correctly by the flattening logic
{2, 1, 0}, // reverse the axes to test that's handled correctly by the flattening logic
{1, 1, 2},
{1, 3, 3},
{0, 21, 6,
@ -559,5 +559,21 @@ TEST(SliceTest, Slice2D_ReverseSubsetOfNegAxes_1) {
true);
}
// test where we provide a subset of axes, we can flatten some dimensions, and we need to skip some input before
// getting to the first value to output.
TEST(SliceTest, Slice5D_SubsetOfAxes_Flatten2Dims_OffsetInput) {
RunSliceTest<float>({1, 2, 2, 2, 2},
{1.f, 2.f, 3.f, 4.f,
5.f, 6.f, 7.f, 8.f,
-1.f, -2.f, -3.f, -4.f,
-5.f, -6.f, -7.f, -8.f},
{0, 1, 1, 0},
{1, 2, std::numeric_limits<int64_t>::max(), 6},
{0, 1, 2, 3},
{},
{1, 1, 1, 2, 2},
{-5.f, -6.f, -7.f, -8.f},
true);
}
} // namespace test
} // namespace onnxruntime