Support negative axis in unsqueeze elimination (#2158)

* Handle negative axes in UnsqueezeElimination.

* Address PR comments

* Trigger rebuild. Some CI builds are unable to be retried

onnxruntime/core/optimizer/unsqueeze_elimination.cc

@@ -12,16 +12,6 @@ using namespace onnxruntime::common;
 namespace onnxruntime {
 
 Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect) const {
-  // Get "axes" attribute. It's a required attribute so can't be null (model loading would fail if it was).
-  const ONNX_NAMESPACE::AttributeProto& attr = *graph_utils::GetNodeAttribute(node, "axes");
-
-  std::vector<int64_t> axes;
-  axes.reserve(attr.ints_size());
-  for (int i = 0; i < attr.ints_size(); i++) {
-    axes.push_back(static_cast<int64_t>(attr.ints(i)));
-  }
-
-  // Generate new dims.
   NodeArg& input_def = *node.MutableInputDefs()[0];
   const auto& tensor_proto = *graph_utils::GetConstantInitializer(graph, input_def.Name());
 
@@ -31,8 +21,24 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect&
     return Status::OK();
   }
 
-  std::vector<int64_t> new_dims(axes.size() + tensor_proto.dims().size(), 0);
+  std::vector<int64_t> axes;
+  if (!graph_utils::GetRepeatedNodeAttributeValues(node, "axes", axes)) {
+    // missing 'axes'. should have failed at model load but just in case...
+    return Status::OK();
+  }
 
+  auto num_axes = axes.size();
+  auto output_rank = num_axes + tensor_proto.dims().size();
+
+  // handle any negative axis values
+  for (auto& axis : axes) {
+    if (axis < 0) {
+      axis += output_rank;
+    }
+  }
+
+  // Generate new dims.
+  std::vector<int64_t> new_dims(output_rank, 0);
   for (int64_t axis : axes) {
     new_dims[axis] = 1;
   }

onnxruntime/test/optimizer/graph_transform_test.cc

@@ -29,6 +29,7 @@
 #include "core/util/math.h"
 #include "test/capturing_sink.h"
 #include "test/framework/test_utils.h"
+#include "test/providers/provider_test_utils.h"
 #include "test/test_environment.h"
 
 #include "gtest/gtest.h"
@@ -262,12 +263,13 @@ TEST(GraphTransformationTests, FuseConvBNNoBias) {
 
 TEST(GraphTransformationTests, FuseConvBNMulAddUnsqueeze) {
   std::vector<std::string> test_models = {"fusion/fuse-conv-bn-mul-add-unsqueeze.onnx",
+                                          "fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx",
                                           "fusion/fuse-conv-bn-mul-add-unsqueeze-no-bias.onnx"};
   for (const auto& model : test_models) {
     string model_uri = MODEL_FOLDER + model;
 
     std::shared_ptr<Model> p_model;
-    ASSERT_TRUE(Model::Load(model_uri, p_model).IsOK());
+    ASSERT_STATUS_OK(Model::Load(model_uri, p_model));
     Graph& graph = p_model->MainGraph();
 
     onnxruntime::GraphTransformerManager graph_transformation_mgr{5};

onnxruntime/test/providers/provider_test_utils.h

@@ -24,6 +24,20 @@
 #include <gsl/gsl>
 #include "core/util/math_cpuonly.h"
 
+// helpers to run a function and check the status, outputting any error if it fails.
+// note: wrapped in do{} while(false) so the _tmp_status variable has limited scope
+#define ASSERT_STATUS_OK(function)                  \
+  do {                                              \
+    auto _tmp_status = function;                    \
+    ASSERT_TRUE(_tmp_status.IsOK()) << _tmp_status; \
+  } while (false)
+
+#define EXPECT_STATUS_OK(function)                  \
+  do {                                              \
+    auto _tmp_status = function;                    \
+    EXPECT_TRUE(_tmp_status.IsOK()) << _tmp_status; \
+  } while (false)
+
 namespace onnxruntime {
 class InferenceSession;
 struct SessionOptions;