From 46eb7187fe0d1a456a0ed268ee3aa30ed8f08aaa Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Tue, 16 Aug 2022 16:52:36 +0800 Subject: [PATCH 1/6] fix(StackOp): fix bug when input number is 128*n+1 fix #8918 --- oneflow/core/functional/impl/array_functor.cpp | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/oneflow/core/functional/impl/array_functor.cpp b/oneflow/core/functional/impl/array_functor.cpp index 13f36692abf..ae780844e9a 100644 --- a/oneflow/core/functional/impl/array_functor.cpp +++ b/oneflow/core/functional/impl/array_functor.cpp @@ -596,7 +596,6 @@ class StackFunctor { int64_t ndims = inputs[0]->ndim(); int64_t stack_dim = dim; stack_dim = JUST(maybe_wrap_dim(stack_dim, ndims + 1)); - if (ninput == 1) { return ExpandDims(inputs[0], dim); } const std::shared_ptr& first_in_shape = inputs[0]->shape(); for (const auto& input : inputs) { for (int i = 0; i < ndims; ++i) { @@ -615,8 +614,12 @@ class StackFunctor { size_t size = (i + kMaxInputCount) < ninput ? kMaxInputCount : ninput - i; TensorTuple partial_inputs(size); for (int j = 0; j < size; ++j) { partial_inputs[j] = inputs[i + j]; } - outputs.emplace_back( - JUST(OpInterpUtil::Dispatch(*ops_.at(size - 1), partial_inputs, attrs))); + if (partial_inputs.size() == 1) { + outputs.emplace_back(JUST(functional::ExpandDims(partial_inputs[0], dim))); + } else { + outputs.emplace_back( + JUST(OpInterpUtil::Dispatch(*ops_.at(size - 1), partial_inputs, attrs))); + } } if (outputs.size() == 1) { return outputs.at(0); } return Concat(outputs, stack_dim); From 611cb3b5a93b97b74a2aab9b1ec5865bca8d5971 Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Tue, 16 Aug 2022 17:08:47 +0800 Subject: [PATCH 2/6] test(Stack): add kMaxInputCount for stack op test --- python/oneflow/test/modules/test_stack.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/python/oneflow/test/modules/test_stack.py b/python/oneflow/test/modules/test_stack.py index cd1f209ce76..bff921552ae 100644 --- a/python/oneflow/test/modules/test_stack.py +++ b/python/oneflow/test/modules/test_stack.py @@ -44,6 +44,13 @@ def test_stack_bool_with_random_data(test_case): out = torch.stack((x, y), dim=random(low=1, high=4).to(int)) return out + @autotest(auto_backward=True, check_graph=True) + def test_stack_kMaxInputCount_inputs(test_case): + kMaxInputCount = 129 + stack_list = [random_tensor(ndim=2, dim0=3, dim1=4) for _ in range(kMaxInputCount)] + out = torch.stack(stack_list, 0) + return out + if __name__ == "__main__": unittest.main() From 6a4815c5d2c1fb301dcad1afc7bcf4e71b6322a0 Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Tue, 16 Aug 2022 17:12:18 +0800 Subject: [PATCH 3/6] format code --- python/oneflow/test/modules/test_stack.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/python/oneflow/test/modules/test_stack.py b/python/oneflow/test/modules/test_stack.py index bff921552ae..7bab6e8b3b4 100644 --- a/python/oneflow/test/modules/test_stack.py +++ b/python/oneflow/test/modules/test_stack.py @@ -47,7 +47,9 @@ def test_stack_bool_with_random_data(test_case): @autotest(auto_backward=True, check_graph=True) def test_stack_kMaxInputCount_inputs(test_case): kMaxInputCount = 129 - stack_list = [random_tensor(ndim=2, dim0=3, dim1=4) for _ in range(kMaxInputCount)] + stack_list = [ + random_tensor(ndim=2, dim0=3, dim1=4) for _ in range(kMaxInputCount) + ] out = torch.stack(stack_list, 0) return out From 39d34c1bf022379468316b26f8d120fe5e660bae Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Tue, 16 Aug 2022 17:14:28 +0800 Subject: [PATCH 4/6] add comment --- oneflow/core/functional/impl/array_functor.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/oneflow/core/functional/impl/array_functor.cpp b/oneflow/core/functional/impl/array_functor.cpp index ae780844e9a..b9e86a6964f 100644 --- a/oneflow/core/functional/impl/array_functor.cpp +++ b/oneflow/core/functional/impl/array_functor.cpp @@ -615,6 +615,7 @@ class StackFunctor { TensorTuple partial_inputs(size); for (int j = 0; j < size; ++j) { partial_inputs[j] = inputs[i + j]; } if (partial_inputs.size() == 1) { + // Use ExpandDims functor for only one input outputs.emplace_back(JUST(functional::ExpandDims(partial_inputs[0], dim))); } else { outputs.emplace_back( From c476b0bacae892ffb4e3e88e576719cd2f90f954 Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Tue, 16 Aug 2022 17:24:16 +0800 Subject: [PATCH 5/6] refine test case --- python/oneflow/test/modules/test_stack.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/oneflow/test/modules/test_stack.py b/python/oneflow/test/modules/test_stack.py index 7bab6e8b3b4..f694a28f5c1 100644 --- a/python/oneflow/test/modules/test_stack.py +++ b/python/oneflow/test/modules/test_stack.py @@ -46,7 +46,7 @@ def test_stack_bool_with_random_data(test_case): @autotest(auto_backward=True, check_graph=True) def test_stack_kMaxInputCount_inputs(test_case): - kMaxInputCount = 129 + kMaxInputCount = 128 + 1 stack_list = [ random_tensor(ndim=2, dim0=3, dim1=4) for _ in range(kMaxInputCount) ] From 67243208658f7d2c1a18f01998a5fe2a1011877b Mon Sep 17 00:00:00 2001 From: wyg1997 Date: Tue, 16 Aug 2022 18:45:30 +0800 Subject: [PATCH 6/6] fix as static analysis --- oneflow/core/functional/impl/array_functor.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/oneflow/core/functional/impl/array_functor.cpp b/oneflow/core/functional/impl/array_functor.cpp index b9e86a6964f..128bccdd56e 100644 --- a/oneflow/core/functional/impl/array_functor.cpp +++ b/oneflow/core/functional/impl/array_functor.cpp @@ -619,7 +619,7 @@ class StackFunctor { outputs.emplace_back(JUST(functional::ExpandDims(partial_inputs[0], dim))); } else { outputs.emplace_back( - JUST(OpInterpUtil::Dispatch(*ops_.at(size - 1), partial_inputs, attrs))); + JUST(OpInterpUtil::Dispatch(*ops_[size - 1], partial_inputs, attrs))); } } if (outputs.size() == 1) { return outputs.at(0); }