fixed bug when broadcast dimensions is negative

Signed-off-by: Chen Xin <[email protected]>
liyuenan2333 · Oct 8, 2022 · 3fed6d6 · 3fed6d6
1 parent a038df2
commit 3fed6d6
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/include/tim/vx/ops/broadcast.h b/include/tim/vx/ops/broadcast.h
@@ -36,10 +36,10 @@ namespace ops {
  *
  * Input:
  * - input.
- * 
+ *
  * Attribute:
  * - shape: the shape which broadcast to.
- * - dimensions(optional): Which dimension in the target shape each dimension 
+ * - dimensions(optional): Which dimension in the target shape each dimension
  *   of the operand shape corresponds to. For BroadcastInDim.
  */
 
@@ -51,7 +51,7 @@ class Broadcast : public BuiltinOp {
 
    protected:
     const std::vector<int32_t> shape_;
-    const std::vector<int32_t> dimensions_;
+    std::vector<int32_t> dimensions_;
 };
 
 }  // namespace ops

diff --git a/src/tim/vx/ops/broadcast.cc b/src/tim/vx/ops/broadcast.cc
@@ -41,6 +41,10 @@ Broadcast::Broadcast(Graph* graph, const std::vector<int32_t>& shape,
   this->impl()->node()->nn_param.expand_broadcast.dimensions_num = dimensions_.size();
   if (dimensions.size() > 0)
   {
+    int dim_num = shape.size();
+    for (uint32_t i = 0; i < dimensions.size(); ++i) {
+      dimensions_[i] += (dimensions[i] < 0 ? dim_num : 0U);
+    }
     this->impl()->node()->nn_param.expand_broadcast.dimensions = (uint32_t*)dimensions_.data();
   } else {
     this->impl()->node()->nn_param.expand_broadcast.dimensions = nullptr;