repeat_interleave support bf16 dtype (#61854) (#61899)

* repeat_interleave support bf16 dtype * support bf16 on cpu

Author: tianhaodongbd

paddle/phi/kernels/cpu/repeat_interleave_grad_kernel.cc

@@ -104,7 +104,8 @@ PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index_grad,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(repeat_interleave_grad,
  CPU,
@@ -113,4 +114,5 @@ PD_REGISTER_KERNEL(repeat_interleave_grad,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}

paddle/phi/kernels/cpu/repeat_interleave_kernel.cc

@@ -25,7 +25,8 @@ PD_REGISTER_KERNEL(repeat_interleave,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index,
  CPU,
@@ -34,4 +35,5 @@ PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/repeat_interleave_grad_kernel.cu

@@ -25,12 +25,14 @@ PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index_grad,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}
 PD_REGISTER_KERNEL(repeat_interleave_grad,
  GPU,
  ALL_LAYOUT,
  phi::RepeatInterleaveGradKernel,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}

paddle/phi/kernels/gpu/repeat_interleave_kernel.cu

@@ -25,7 +25,8 @@ PD_REGISTER_KERNEL(repeat_interleave,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}
 
 PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index,
  GPU,
@@ -34,4 +35,5 @@ PD_REGISTER_KERNEL(repeat_interleave_with_tensor_index,
  float,
  double,
  int,
- int64_t) {}
+ int64_t,
+ phi::dtype::bfloat16) {}

test/legacy_test/test_repeat_interleave_op.py

@@ -252,6 +252,25 @@ def test_dygraph_api(self):
  expect_out = np.repeat(input_x, index, axis=None)
  np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
 
+ # case input dtype is bfloat16
+ input_x = np.array([[1, 2, 1], [1, 2, 3]]).astype('uint16')
+
+ with base.dygraph.guard():
+ x = paddle.to_tensor(input_x)
+ index = paddle.to_tensor(index_x)
+ z = paddle.repeat_interleave(x, index, None)
+ np_z = z.numpy()
+ expect_out = np.repeat(input_x, index_x, axis=None)
+ np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
+
+ with base.dygraph.guard():
+ x = paddle.to_tensor(input_x)
+ index = 2
+ z = paddle.repeat_interleave(x, index, None)
+ np_z = z.numpy()
+ expect_out = np.repeat(input_x, index, axis=None)
+ np.testing.assert_allclose(expect_out, np_z, rtol=1e-05)
+
  # case 1:
  with base.dygraph.guard():
  x = base.dygraph.to_variable(self.data_x)