[0-size Tensor Job2 No.16、85] Add 0-size Tensor support for paddle.gather_nd [fluid_ops]

paddle/phi/kernels/cpu/gather_nd_kernel.cc

@@ -17,6 +17,7 @@
 #include "paddle/phi/backends/cpu/cpu_context.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/gather.h"
+#include "paddle/phi/kernels/tile_kernel.h"
 
 namespace phi {
 
@@ -27,6 +28,21 @@ void GatherNdKernel(const Context &dev_ctx,
  DenseTensor *out) {
  dev_ctx.template Alloc<T>(out);
  if (x.numel() == 0 || out->numel() == 0) return;
+ // The result dims is
+ // Index.shape[:-1] + X.shape[Index.shape[-1]:]
+ // If the last dimension of index is 0, set it to 1 and tile x.
+ auto index_dims = index.dims();
+ std::vector<int64_t> out_dims;
+ if (index_dims[index_dims.size() - 1] == 0) {
+ for (int i = 0; i < index_dims.size() - 1; ++i) {
+ out_dims.emplace_back(index_dims[i]);
+ }
+ for (int i = 0; i < x.dims().size(); ++i) {
+ out_dims.emplace_back(1);
+ }
+ phi::TileKernel<T, Context>(dev_ctx, x, phi::IntArray(out_dims), out);
+ return;
+ }
  if (index.dims()[0] == 0 && index.numel() == 0) return;
  auto index_type = index.dtype();
  bool index_type_match =

paddle/phi/kernels/cpu/tile_grad_kernel.cc

@@ -26,6 +26,9 @@ PD_REGISTER_KERNEL(tile_grad,
  float,
  double,
  int,
+ int64_t,
+ int8_t,
+ int16_t,
+ uint8_t,
  phi::dtype::complex<float>,
- phi::dtype::complex<double>,
- int64_t) {}
+ phi::dtype::complex<double>) {}

paddle/phi/kernels/cpu/tile_kernel.cc

@@ -27,6 +27,9 @@ PD_REGISTER_KERNEL(tile,
  double,
  int,
  int64_t,
+ int8_t,
+ int16_t,
+ uint8_t,
  phi::dtype::float16,
  phi::dtype::complex<float>,
  phi::dtype::complex<double>) {}

paddle/phi/kernels/funcs/eigen/broadcast.cu

@@ -14,8 +14,9 @@ limitations under the License. */
 #include "paddle/phi/common/bfloat16.h"
 #include "paddle/phi/common/complex.h"
 #include "paddle/phi/common/float16.h"
+#include "paddle/phi/common/float8_e4m3fn.h"
+#include "paddle/phi/common/float8_e5m2.h"
 #include "paddle/phi/kernels/funcs/eigen/eigen_function.h"
-
 namespace phi {
 namespace funcs {
 
@@ -93,6 +94,11 @@ INSTANTIATION(EigenBroadcastGrad, dtype::complex<float>);
 INSTANTIATION(EigenBroadcastGrad, dtype::complex<double>);
 INSTANTIATION(EigenBroadcastGrad, int);
 INSTANTIATION(EigenBroadcastGrad, int64_t);
+INSTANTIATION(EigenBroadcastGrad, int8_t);
+INSTANTIATION(EigenBroadcastGrad, uint8_t);
+INSTANTIATION(EigenBroadcastGrad, int16_t);
+INSTANTIATION(EigenBroadcastGrad, phi::dtype::float8_e4m3fn);
+INSTANTIATION(EigenBroadcastGrad, phi::dtype::float8_e5m2);
 template struct EigenBroadcastGrad<Eigen::GpuDevice, float, 0>;
 template struct EigenBroadcastGrad<Eigen::GpuDevice, dtype::float16, 0>;
 template struct EigenBroadcastGrad<Eigen::GpuDevice, double, 0>;

paddle/phi/kernels/gpu/gather_nd_kernel.cu

@@ -18,6 +18,7 @@
 #include "paddle/phi/common/bfloat16.h"
 #include "paddle/phi/core/kernel_registry.h"
 #include "paddle/phi/kernels/funcs/gather.cu.h"
+#include "paddle/phi/kernels/tile_kernel.h"
 
 namespace phi {
 
@@ -28,6 +29,21 @@ void GatherNdKernel(const Context &dev_ctx,
  DenseTensor *out) {
  dev_ctx.template Alloc<T>(out);
  if (x.numel() == 0 || out->numel() == 0) return;
+ // The result dims is
+ // Index.shape[:-1] + X.shape[Index.shape[-1]:]
+ // If the last dimension of index is 0, set it to 1 and tile x.
+ auto index_dims = index.dims();
+ std::vector<int64_t> out_dims;
+ if (index_dims[index_dims.size() - 1] == 0) {
+ for (int i = 0; i < index_dims.size() - 1; ++i) {
+ out_dims.emplace_back(index_dims[i]);
+ }
+ for (int i = 0; i < x.dims().size(); ++i) {
+ out_dims.emplace_back(1);
+ }
+ phi::TileKernel<T, Context>(dev_ctx, x, phi::IntArray(out_dims), out);
+ return;
+ }
  if (index.dims()[0] == 0 && index.numel() == 0) return;
  const auto &index_type = index.dtype();
  bool index_type_match =

paddle/phi/kernels/gpu/tile_grad_kernel.cu

@@ -27,7 +27,12 @@ PD_REGISTER_KERNEL(tile_grad,
  double,
  int,
  int64_t,
+ int8_t,
+ int16_t,
+ uint8_t,
  phi::dtype::float16,
+ phi::dtype::bfloat16,
+ phi::dtype::float8_e4m3fn,
+ phi::dtype::float8_e5m2,
  phi::dtype::complex<float>,
- phi::dtype::complex<double>,
- phi::dtype::bfloat16) {}
+ phi::dtype::complex<double>) {}

paddle/phi/kernels/gpu/tile_kernel.cu

@@ -107,6 +107,9 @@ PD_REGISTER_KERNEL(tile,
  double,
  int,
  int64_t,
+ int8_t,
+ int16_t,
+ uint8_t,
  phi::dtype::float16,
  phi::dtype::bfloat16,
  phi::dtype::float8_e4m3fn,

paddle/phi/kernels/xpu/gather_nd_kernel.cc

@@ -16,6 +16,7 @@
 
 #include "paddle/phi/backends/xpu/enforce_xpu.h"
 #include "paddle/phi/core/kernel_registry.h"
+#include "paddle/phi/kernels/tile_kernel.h"
 
 namespace phi {
 
@@ -28,6 +29,21 @@ void GatherNdKernel(const Context &dev_ctx,
  dev_ctx.template Alloc<T>(out);
 
  if (x.numel() == 0 || out->numel() == 0) return;
+ // The result dims is
+ // Index.shape[:-1] + X.shape[Index.shape[-1]:]
+ // If the last dimension of index is 0, set it to 1 and tile x.
+ auto index_dims = index.dims();
+ std::vector<int64_t> out_dims;
+ if (index_dims[index_dims.size() - 1] == 0) {
+ for (int i = 0; i < index_dims.size() - 1; ++i) {
+ out_dims.emplace_back(index_dims[i]);
+ }
+ for (int i = 0; i < x.dims().size(); ++i) {
+ out_dims.emplace_back(1);
+ }
+ phi::TileKernel<T, Context>(dev_ctx, x, phi::IntArray(out_dims), out);
+ return;
+ }
  if (index.dims()[0] == 0 && index.numel() == 0) return;
 
  if (index.numel() == 0) {

test/legacy_test/test_gather_nd_op.py

@@ -721,6 +721,42 @@ def test_imperative(self):
  paddle.enable_static()
 
 
+class TestGatherNdOp_ZeroSize(OpTest):
+ def setUp(self):
+ self.op_type = "gather_nd"
+ self.python_api = paddle.gather_nd
+ self.public_python_api = paddle.gather_nd
+ xnp = np.random.random([10, 20])
+ index = np.random.random([0]).astype("int32")
+ output = xnp[tuple(index.T)]
+
+ self.inputs = {'X': xnp, 'Index': index}
+ self.outputs = {'Out': output}
+
+ def test_check_output(self):
+ self.check_output(check_pir=True)
+
+ def test_check_grad(self):
+ self.check_grad(
+ ['X'],
+ 'Out',
+ check_pir=True,
+ )
+
+
+class TestGatherNdOp_ZeroSize2(TestGatherNdOp_ZeroSize):
+ def setUp(self):
+ self.op_type = "gather_nd"
+ self.python_api = paddle.gather_nd
+ self.public_python_api = paddle.gather_nd
+ xnp = np.random.random([10, 20])
+ index = np.random.random([2, 0]).astype("int32")
+ output = np.tile(xnp, [2, 1, 1])
+
+ self.inputs = {'X': xnp, 'Index': index}
+ self.outputs = {'Out': output}
+
+
 if __name__ == "__main__":
  paddle.enable_static()
  unittest.main()