[GitHub] [tvm] comaniac commented on a change in pull request #7951: [Frontend][Tensorflow]add batch_dim support for gatherV2

[GitBox <gi...@apache.org>]

comaniac commented on a change in pull request #7951:
URL: https://github.com/apache/tvm/pull/7951#discussion_r624223779



##########
File path: include/tvm/topi/transform.h
##########
@@ -854,46 +857,99 @@ inline Tensor sequence_mask(const Tensor& data, const Tensor& valid_length, doub
  *
  * \return A Tensor whose op member is the take operation
  */
-inline Tensor take(const Tensor& a, const Tensor& indices, int axis, std::string mode = "clip",
-                   std::string name = "T_take", std::string tag = kInjective) {
+inline Tensor take(const Tensor& a, const Tensor& indices, int batch_dims, int axis,
+                   std::string mode = "clip", std::string name = "T_take",
+                   std::string tag = kInjective) {
   if (axis < 0) {
     axis += static_cast<int>(a->shape.size());
   }
   ICHECK_GE(axis, 0) << "axis out of bounds";
   ICHECK_LT(axis, a->shape.size()) << "axis out of bounds";
   auto axis_dim = a->shape[axis];
-
   int indices_len = static_cast<int>(indices->shape.size());
-  Array<PrimExpr> out_shape;
-  for (size_t i = 0; i < a->shape.size(); ++i) {
-    if (axis == static_cast<int>(i)) {
-      for (size_t j = 0; j < indices->shape.size(); ++j) {
-        out_shape.push_back(indices->shape[j]);
-      }
-    } else {
-      out_shape.push_back(a->shape[i]);
+
+  int batch_dims_ = batch_dims;
+  if (batch_dims_ != 0) {
+    ICHECK_GE(batch_dims_, -static_cast<int>(indices->shape.size())) << "batch_dims out of bounds";
+    ICHECK_LE(batch_dims_, indices->shape.size()) << "batch_dims out of bounds";
+
+    if (batch_dims_ < 0) {
+      batch_dims_ = indices->shape.size() + batch_dims_;
     }
+
+    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";

Review comment:
       ```suggestion
       ICHECK_LE(batch_dims_, axis) << "batch_dims must be less than or equal to axis";
   ```

##########
File path: src/topi/transform.cc
##########
@@ -87,13 +87,15 @@ TVM_REGISTER_GLOBAL("topi.layout_transform").set_body([](TVMArgs args, TVMRetVal
 });
 
 TVM_REGISTER_GLOBAL("topi.take").set_body([](TVMArgs args, TVMRetValue* rv) {
-  if (args.size() == 3) {
-    std::string mode = args[2];
-    *rv = take(args[0], args[1], mode);
-  } else {
-    int axis = args[2];
+  if (args.size() == 4) {
     std::string mode = args[3];
-    *rv = take(args[0], args[1], axis, mode);
+    int batch_dims = args[2];
+    *rv = take(args[0], args[1], batch_dims, mode);
+  } else {
+    int batch_dims = args[2];
+    int axis = args[3];
+    std::string mode = args[4];

Review comment:
       So the nargs is expected to be 5? Please add a CHECK.




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org