[NVIDIA] Fix Concat/VariadicSplit layers with negative axis (#476)

* Fix Concat layer with negative axis

* Fix negative axis for VariadicSplit

Author: nkogteva

modules/nvidia_plugin/src/ops/concat.cpp

@@ -29,7 +29,10 @@ ConcatOp::ConcatOp(const CreationContext& context,
     Expects(num_inputs_ == GetInputIds().size());
     Expects(GetOutputIds().size() == 1);
     const auto& outputShape = concatOp.get_output_shape(0);
-    const int64_t axis = concatOp.get_axis();
+    int64_t axis = concatOp.get_axis();
+    if (axis < 0) {
+        axis += static_cast<int64_t>(concatOp.get_input_partial_shape(0).rank().get_length());
+    }
     Expects(axis >= 0 && axis < outputShape.size());
     auto num_chunks =
         std::accumulate(outputShape.begin(), outputShape.begin() + axis + 1, 1, std::multiplies<size_t>());

modules/nvidia_plugin/src/ops/variadic_split.cpp

@@ -92,7 +92,10 @@ VariadicSplitOp::VariadicSplitOp(const CreationContext& context,
     const auto element_type = input_element_type;
 
     const auto& data_shape = variadic_split_node->get_input_shape(0);
-    const int64_t axis = *axis_node->get_data_ptr<int64_t>();
+    int64_t axis = *axis_node->get_data_ptr<int64_t>();
+    if (axis < 0) {
+        axis += static_cast<int64_t>(variadic_split_node->get_input_partial_shape(0).rank().get_length());
+    }
     Expects(axis >= 0 && axis < data_shape.size());
     const size_t orig_axis_size = data_shape[axis];
 

modules/nvidia_plugin/tests/functional/shared_tests_instances/single_layer_tests/concat.cpp

@@ -13,7 +13,7 @@ using namespace LayerTestsDefinitions;
 
 namespace {
 
-std::vector<int> axes = {0, 1, 2, 3};
+std::vector<int> axes = {-3, -2, -1, 0, 1, 2, 3};
 std::vector<std::vector<std::vector<size_t>>> inShapes = {
     {{10, 10, 10, 10}},
     {{10, 10, 10, 10}, {10, 10, 10, 10}},

modules/nvidia_plugin/tests/functional/shared_tests_instances/single_layer_tests/variadic_split.cpp

@@ -30,7 +30,7 @@ const std::vector<std::vector<size_t>> smoke_numSplits = {
 INSTANTIATE_TEST_CASE_P(smoke_NumSplitsCheck,
                         VariadicSplitLayerTest,
                         ::testing::Combine(::testing::ValuesIn(smoke_numSplits),
-                                           ::testing::Values(0, 1, 2, 3),
+                                           ::testing::Values(-3, -2, -1, 0, 1, 2, 3),
                                            ::testing::ValuesIn(netPrecisions),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),
                                            ::testing::Values(InferenceEngine::Precision::UNSPECIFIED),