fix as_complex,as_real strided kernel bug (#61391)

wanghuancoder · web-flow · commit 2d24faec8f7f · 2024-02-01T10:01:11.000+08:00
diff --git a/paddle/phi/kernels/stride/as_complex_kernel.cc b/paddle/phi/kernels/stride/as_complex_kernel.cc
@@ -23,7 +23,22 @@ template <typename T, typename Context>
 void AsComplexStridedKernel(const Context& dev_ctx,
                             const DenseTensor& x,
                             DenseTensor* out) {
-  out->set_strides(DenseTensorMeta::calc_strides(out->dims()));
+  PADDLE_ENFORCE_EQ(
+      x.strides()[x.strides().size() - 1],
+      1,
+      phi::errors::InvalidArgument(
+          "Expected the stride of last dimension of input(X) to be 1."
+          "But received %d. This means that the last dimension of the"
+          "Tensor(x) is not continuous and cannot be as_complex directly."
+          "You can call x.contiguous() to make the Tensor(x) contiguous first.",
+          x.strides()[x.strides().size() - 1]));
+
+  phi::DDim out_strides(x.strides().Get(), x.strides().size() - 1);
+  for (int i = 0; i < out_strides.size(); ++i) {
+    out_strides[i] /= 2;
+  }
+  out->set_strides(out_strides);
+
   if (x.dtype() == DataType::FLOAT32) {
     out->set_type(DataType::COMPLEX64);
   } else if (x.dtype() == DataType::FLOAT64) {
diff --git a/paddle/phi/kernels/stride/as_real_kernel.cc b/paddle/phi/kernels/stride/as_real_kernel.cc
@@ -21,7 +21,13 @@ template <typename T, typename Context>
 void AsRealStridedKernel(const Context& dev_ctx,
                          const DenseTensor& x,
                          DenseTensor* out) {
-  out->set_strides(DenseTensorMeta::calc_strides(out->dims()));
+  auto out_stride_v = common::vectorize(x.strides());
+  for (size_t i = 0; i < out_stride_v.size(); ++i) {
+    out_stride_v[i] *= 2;
+  }
+  out_stride_v.push_back(1);
+  out->set_strides(common::make_ddim(out_stride_v));
+
   if (x.dtype() == DataType::COMPLEX64) {
     out->set_type(DataType::FLOAT32);
   } else if (x.dtype() == DataType::COMPLEX128) {
