add check for num of arguments for topi.take

apache · Apr 30, 2021 · bfe7ddf · bfe7ddf
1 parent 1423c67
commit bfe7ddf
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/include/tvm/topi/transform.h b/include/tvm/topi/transform.h
@@ -878,7 +878,7 @@ inline Tensor take(const Tensor& a, const Tensor& indices, int batch_dims, int a
     }
 
     ICHECK_LT(batch_dims_, a->shape.size()) << "batch_dims out of bounds";
-    ICHECK_GE(axis, batch_dims_) << "batch_dims must be less than or equal to axis";
+    ICHECK_LE(batch_dims_, axis) << "batch_dims must be less than or equal to axis";
     for (int i = 0; i < batch_dims_; ++i) {
       auto addr1 = a->shape[i];
       auto addr2 = indices->shape[i];

diff --git a/src/topi/transform.cc b/src/topi/transform.cc
@@ -92,6 +92,7 @@ TVM_REGISTER_GLOBAL("topi.take").set_body([](TVMArgs args, TVMRetValue* rv) {
     int batch_dims = args[2];
     *rv = take(args[0], args[1], batch_dims, mode);
   } else {
+    ICHECK_EQ(args.size(), 5) << "topi.take expects 4 or 5 arguments";
     int batch_dims = args[2];
     int axis = args[3];
     std::string mode = args[4];