check

paddle/phi/kernels/cpu/gather_tree_kernel.cc

@@ -49,6 +49,15 @@ void GatherTreeKernel(const Context &dev_ctx,
       out_data[idx] = ids_data[idx];
       auto parent = parents_data[idx];
       for (int step = max_length - 2; step >= 0; step--) {
+        PADDLE_ENFORCE_LT(
+            parent,
+            beam_size,
+            paddle::platform::errors::InvalidArgument(
+                "The parents must be less than beam size, but recieved"
+                "parents %d is greater than or equal to beam size %d. ",
+                parent,
+                beam_size));
+
         idx = step * batch_size * beam_size + batch * beam_size;
         out_data[idx + beam] = ids_data[idx + parent];
         parent = parents_data[idx + parent];

paddle/phi/kernels/gather_tree_kernel.h

@@ -14,6 +14,7 @@
 
 #pragma once
 
+#include "paddle/fluid/platform/enforce.h"
 #include "paddle/phi/core/dense_tensor.h"
 
 namespace phi {

paddle/phi/kernels/gpu/gather_tree_kernel.cu

@@ -35,6 +35,9 @@ __global__ void GatherTree(const T *ids_data,
     out_data[idx] = ids_data[idx];
     auto parent = parents_data[idx];
     for (int step = max_length - 2; step >= 0; step--) {
+      assert(parent < beam_size &&
+             "The parents of gather_tree op must be less than beam size. ");
+
       idx = step * batch_size * beam_size + batch * beam_size;
       out_data[idx + beam] = ids_data[idx + parent];
       parent = parents_data[idx + parent];

python/paddle/nn/functional/extension.py

@@ -303,6 +303,13 @@ def gather_tree(ids, parents):
             # [[[2, 2], [1, 6]], [[3, 3], [6, 1]], [[0, 1], [9, 0]]]
 
     """
+    if ids.ndim != 3:
+        raise ValueError(
+            "The input ids must be a 3D tensor with shape [length, batch_size, beam_size]"
+        )
+    if ids.ndim != parents.ndim:
+        raise ValueError("The ids's shape must be the same as parents' shape. ")
+
     if in_dygraph_mode():
         return _C_ops.gather_tree(ids, parents)
     else: