tool(core/util/): added tool to unsqueeze dimensions at a specified index

Signed-off-by: Abhiram Iyer <[email protected]>

Signed-off-by: Abhiram Iyer <[email protected]>

Author: abhi-iyer

core/util/trt_util.cpp

@@ -106,6 +106,32 @@ nvinfer1::Dims unpadDims(const nvinfer1::Dims& d) {
     return dims;
 }
 
+nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos) {
+    // acceptable range for pos is [0, d.nbDims]
+    TRTORCH_ASSERT(pos >= 0 && pos <= d.nbDims, "ERROR: Index to unsqueeze is out of bounds.");
+    
+    nvinfer1::Dims dims;
+
+    int i = 0;
+    int j = 0;
+
+    while (i <= d.nbDims) {
+        if (j != pos) {
+            dims.d[j] = d.d[i];
+            i++;
+        } else {
+            // add new dimension at pos
+            dims.d[j] = 1;
+        }
+
+        j++;
+    }
+
+    dims.nbDims = d.nbDims+1;
+
+    return dims;
+}
+
 std::vector<int64_t> toVec(nvinfer1::Dims d) {
     std::vector<int64_t> dims;
     for (int i = 0; i < d.nbDims; i++) {

core/util/trt_util.h

@@ -80,6 +80,7 @@ int64_t volume(const nvinfer1::Dims& d);
 nvinfer1::Dims toDimsPad(c10::IntArrayRef l, uint64_t pad_to);
 nvinfer1::Dims toDimsPad(c10::List<int64_t> l, uint64_t pad_to);
 nvinfer1::Dims unpadDims(const nvinfer1::Dims& d);
+nvinfer1::Dims unsqueezeDims(const nvinfer1::Dims& d, int pos);
 nvinfer1::Dims toDims(c10::IntArrayRef l);
 nvinfer1::Dims toDims(c10::List<int64_t> l);
 nvinfer1::DimsHW toDimsHW(c10::List<int64_t> l);