From aefa466334f43a89c0821f74aef7f354f799e356 Mon Sep 17 00:00:00 2001 From: Changming Sun Date: Wed, 1 Apr 2020 16:20:14 -0700 Subject: [PATCH] Allow zero in split op (#3389) Allow zero in split op (A change in onnx 1.7 without bumping up the op version) --- onnxruntime/core/providers/cpu/sequence/sequence_ops.cc | 9 +++------ onnxruntime/core/providers/cpu/tensor/split.h | 2 +- onnxruntime/test/providers/cpu/tensor/split_op_test.cc | 4 ++-- 3 files changed, 6 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/providers/cpu/sequence/sequence_ops.cc b/onnxruntime/core/providers/cpu/sequence/sequence_ops.cc index 42af8e607b..e9e15b8b99 100644 --- a/onnxruntime/core/providers/cpu/sequence/sequence_ops.cc +++ b/onnxruntime/core/providers/cpu/sequence/sequence_ops.cc @@ -472,17 +472,14 @@ Status SplitToSequence::ComputeImpl(OpKernelContext& context, const Tensor& inpu is_split_input_scalar = true; } else { GetSplitSizesInput(*p_split_input, split_sizes); - ORT_ENFORCE(std::all_of(split_sizes.cbegin(), split_sizes.cend(), [](int64_t value) { return value > 0; }), - "Invalid value in 'split' input. All values must be > 0"); + ORT_ENFORCE(std::all_of(split_sizes.cbegin(), split_sizes.cend(), [](int64_t value) { return value >= 0; }), + "Invalid value in 'split' input. All values must be >= 0"); } } // Keep the split dimension or not. Default 1, which means we keep split dimension. // If input 'split' is specified, this attribute is ignored. - bool use_keep_dims = false; - if (split_sizes.empty()) { - use_keep_dims = true; - } + bool use_keep_dims = split_sizes.empty(); ORT_RETURN_IF_ERROR(PrepareForCompute(input_shape, split_scalar, diff --git a/onnxruntime/core/providers/cpu/tensor/split.h b/onnxruntime/core/providers/cpu/tensor/split.h index 51b13b51be..d2e51230a7 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.h +++ b/onnxruntime/core/providers/cpu/tensor/split.h @@ -18,7 +18,7 @@ class SplitBase { // optional if (info.GetAttrs("split", split_sizes_).IsOK()) { split_size_sum_ = std::accumulate(split_sizes_.cbegin(), split_sizes_.cend(), 0LL); - ORT_ENFORCE(std::all_of(split_sizes_.cbegin(), split_sizes_.cend(), [](int64_t value) { return value > 0; }), + ORT_ENFORCE(std::all_of(split_sizes_.cbegin(), split_sizes_.cend(), [](int64_t value) { return value >= 0; }), "Invalid value in 'split' attribute. All values must be > 0"); } } diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index 1ff8c4d79d..7b1a582e2f 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -429,7 +429,7 @@ TEST(SplitOperatorTest, SplitAttributeSumTooSmall) { } TEST(SplitOperatorTest, InvalidValueInSplitAttribute) { - const int64_t axis = 0; + const int64_t axis = -1; std::vector outputs; // input shape and data @@ -443,7 +443,7 @@ TEST(SplitOperatorTest, InvalidValueInSplitAttribute) { outputs.push_back({{1, 2}, {1.f, 2.f}}); outputs.push_back({{3, 2}, {3.f, 4.f, 5.f, 6.f, 7.f, 8.f}}); - RunTest(axis, splits, input, outputs, false, true, "Invalid value in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, splits, input, outputs, false, true, "in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM } /*