PaddlePaddle
diff --git a/‎paddle/fluid/operators/beam_search_op.cc
Copy file name to clipboard
-159Lines changed: 0 additions & 159 deletions b/‎paddle/fluid/operators/beam_search_op.cc
Copy file name to clipboard
-159Lines changed: 0 additions & 159 deletions
diff --git a/‎paddle/fluid/operators/beam_search_op.cu.cc
Copy file name to clipboard
-28Lines changed: 0 additions & 28 deletions b/‎paddle/fluid/operators/beam_search_op.cu.cc
Copy file name to clipboard
-28Lines changed: 0 additions & 28 deletions
diff --git a/‎paddle/fluid/operators/beam_search_op.h
Copy file name to clipboard
-77Lines changed: 0 additions & 77 deletions b/‎paddle/fluid/operators/beam_search_op.h
Copy file name to clipboard
-77Lines changed: 0 additions & 77 deletions
diff --git a/‎paddle/fluid/operators/beam_search_op_xpu.cc
Copy file name to clipboard
-29Lines changed: 0 additions & 29 deletions b/‎paddle/fluid/operators/beam_search_op_xpu.cc
Copy file name to clipboard
-29Lines changed: 0 additions & 29 deletions
diff --git a/‎paddle/phi/infermeta/multiary.cc
Copy file name to clipboard
+20Lines changed: 20 additions & 0 deletions b/‎paddle/phi/infermeta/multiary.cc
Copy file name to clipboard
+20Lines changed: 20 additions & 0 deletions
diff --git a/‎paddle/phi/infermeta/multiary.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎paddle/phi/infermeta/multiary.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/cpu/beam_search_kernel.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/phi/kernels/cpu/beam_search_kernel.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/phi/kernels/gpu/beam_search_kernel.cu
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/phi/kernels/gpu/beam_search_kernel.cu
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
@@ -894,6 +894,26 @@ void BilinearInferMeta(const MetaTensor& x,
   out->set_dtype(x.dtype());
 }
 
+void BeamSearchInferMeta(const MetaTensor& pre_ids,
+                         const MetaTensor& pre_scores,
+                         const MetaTensor& ids,
+                         const MetaTensor& scores,
+                         int level,
+                         int beam_size,
+                         int end_id,
+                         bool is_accumulated,
+                         MetaTensor* selected_ids,
+                         MetaTensor* selected_scores,
+                         MetaTensor* parent_idx) {
+  const auto& id_dims = pre_ids.dims();
+  selected_scores->set_dims(pre_scores.dims());
+  selected_ids->set_dims(id_dims);
+  parent_idx->set_dims({id_dims[0]});
+  selected_scores->set_dtype(pre_scores.dtype());
+  selected_ids->set_dtype(pre_ids.dtype());
+  parent_idx->set_dtype(pre_ids.dtype());
+}
+
 void BroadcastTensorsInferMeta(const std::vector<const MetaTensor*>& x,
                                std::vector<MetaTensor*> out) {
   int target_rank = 0;
 
@@ -213,6 +213,18 @@ void BatchNormInferInferMeta(const MetaTensor& x,
                              MetaTensor* variance_out,
                              MetaConfig config = MetaConfig());
 
+void BeamSearchInferMeta(const MetaTensor& pre_ids,
+                         const MetaTensor& pre_scores,
+                         const MetaTensor& ids,
+                         const MetaTensor& scores,
+                         int level,
+                         int beam_size,
+                         int end_id,
+                         bool is_accumulated,
+                         MetaTensor* selected_ids,
+                         MetaTensor* selected_scores,
+                         MetaTensor* parent_idx);
+
 void BilinearInferMeta(const MetaTensor& x,
                        const MetaTensor& y,
                        const MetaTensor& weight,
 
@@ -0,0 +1,25 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/beam_search_kernel_impl.h"
+
+PD_REGISTER_KERNEL(beam_search,
+                   CPU,
+                   ALL_LAYOUT,
+                   phi::BeamSearchOpKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}
@@ -0,0 +1,25 @@
+// Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/impl/beam_search_kernel_impl.h"
+
+PD_REGISTER_KERNEL(beam_search,
+                   GPU,
+                   ALL_LAYOUT,
+                   phi::BeamSearchOpKernel,
+                   float,
+                   double,
+                   int,
+                   int64_t) {}