Fix

Author: co63oc

paddle/phi/kernels/cpu/roll_grad_kernel.cc

@@ -27,6 +27,10 @@ void RollGradKernel(const Context& dev_ctx,
                     const IntArray& shifts,
                     const std::vector<int64_t>& axis,
                     DenseTensor* x_grad) {
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   std::vector<T> out_vec;
   phi::TensorToVector(out_grad, dev_ctx, &out_vec);
 

paddle/phi/kernels/gpu/roll_grad_kernel.cu

@@ -30,6 +30,10 @@ void RollGradKernel(const Context& dev_ctx,
                     const std::vector<int64_t>& axis,
                     DenseTensor* x_grad) {
   auto* out_grad_data = out_grad.data<T>();
+  if (x_grad && x_grad->numel() == 0) {
+    dev_ctx.template Alloc<T>(x_grad);
+    return;
+  }
   T* x_grad_data = dev_ctx.template Alloc<T>(x_grad);
 
   auto shifts_data = shifts.GetData();

test/fft/test_fft.py

@@ -1934,5 +1934,47 @@ def test_ifftshift(self):
             )
 
 
+def create_test_class(op_type, dtype, shape, axis):
+    class Cls(unittest.TestCase):
+        def test_zero_size(self):
+            paddle.disable_static()
+            import scipy  # noqa: F401
+
+            numpy_tensor_1 = np.random.rand(*shape).astype(dtype)
+            paddle_x = paddle.to_tensor(numpy_tensor_1)
+            paddle_x.stop_gradient = False
+
+            paddle_api = eval(f"paddle.fft.{op_type}")
+            paddle_out = paddle_api(
+                paddle_x, axes=axis
+            )  # Here the parameter is axes
+            numpy_api = eval(f"scipy.fft.{op_type}")
+            numpy_out = numpy_api(numpy_tensor_1, axes=axis)
+
+            np.testing.assert_allclose(
+                paddle_out.numpy(),
+                numpy_out,
+                1e-2,
+                1e-2,
+            )
+            np.testing.assert_allclose(
+                paddle_out.shape,
+                numpy_out.shape,
+            )
+
+    cls_name = f"{op_type}{dtype}_0SizeTest"
+    Cls.__name__ = cls_name
+    globals()[cls_name] = Cls
+
+
+op_list = [
+    ("fftshift", ("complex64", "complex128"), (0, -1)),
+    ("ifftshift", ("complex64", "complex128"), (0, -1)),
+]
+
+for op in op_list:
+    create_test_class(op[0], op[1][0], [3, 4, 0], op[2][0])
+    create_test_class(op[0], op[1][1], [3, 4, 0, 3, 4], op[2][1])
+
 if __name__ == '__main__':
     unittest.main()