address comments

Author: vtavana

CHANGELOG.md

@@ -29,7 +29,7 @@ This release achieves 100% compliance with Python Array API specification (revis
 * Removed `einsum_call` keyword from `dpnp.einsum_path` signature [#2421](https://github.com/IntelPython/dpnp/pull/2421)
 * Changed `"max dimensions"` to `None` in array API capabilities [#2432](https://github.com/IntelPython/dpnp/pull/2432)
 * Updated kernel header `i0.hpp` to expose `cyl_bessel_i0` function depending on build target [#2440](https://github.com/IntelPython/dpnp/pull/2440)
-* Updated FFT module to make input array Hermitian before calling complex-to-real FFT [#2444](https://github.com/IntelPython/dpnp/pull/2444)
+* Updated FFT module to ensure an input array is Hermitian before calling complex-to-real FFT [#2444](https://github.com/IntelPython/dpnp/pull/2444)
 
 ### Fixed
 

dpnp/fft/dpnp_utils_fft.py

@@ -285,30 +285,12 @@ def _copy_array(x, complex_input):
             dtype = map_dtype_to_device(dpnp.float64, x.sycl_device)
 
     if copy_flag:
-        x = _copy_kernel(x, dtype)
+        x = x.astype(dtype, copy=True)
 
     # if copying is done, FFT can be in-place (copy_flag = in_place flag)
     return x, copy_flag
 
 
-def _copy_kernel(x, dtype):
-    x_copy = dpnp.empty_like(x, dtype=dtype, order="C")
-
-    exec_q = x.sycl_queue
-    _manager = dpu.SequentialOrderManager[exec_q]
-    dep_evs = _manager.submitted_events
-
-    ht_copy_ev, copy_ev = ti._copy_usm_ndarray_into_usm_ndarray(
-        src=dpnp.get_usm_ndarray(x),
-        dst=x_copy.get_array(),
-        sycl_queue=exec_q,
-        depends=dep_evs,
-    )
-    _manager.add_event_pair(ht_copy_ev, copy_ev)
-
-    return x_copy
-
-
 def _extract_axes_chunk(a, s, chunk_size=3):
     """
     Classify the first input into a list of lists with each list containing
@@ -438,9 +420,9 @@ def _fft(a, norm, out, forward, in_place, c2c, axes, batch_fft=True):
     return result
 
 
-def _make_array_hermitian(a, n, copy_needed):
+def _make_array_hermitian(a, n, axis, copy_needed):
     """
-    For `dpnp.fft.irfft`, the input array should be Hermitian. If it is not,
+    For complex-to-real FFT, the input array should be Hermitian. If it is not,
     the behavior is undefined. This function makes necessary changes to make
     sure the given array is Hermitian.
 
@@ -449,6 +431,7 @@ def _make_array_hermitian(a, n, copy_needed):
     `_truncate_or_pad`, so the array has enough length.
     """
 
+    a = dpnp.moveaxis(a, axis, 0)
     length_is_even = n % 2 == 0
     hermitian = dpnp.all(a[0].imag == 0)
     assert n is not None
@@ -463,14 +446,14 @@ def _make_array_hermitian(a, n, copy_needed):
 
     if not hermitian:
         if copy_needed:
-            a = _copy_kernel(a, a.dtype)
+            a = a.astype(a.dtype, copy=True)
 
         a[0].imag = 0
         if length_is_even:
             f_ny = n // 2
             a[f_ny].imag = 0
 
-    return a
+    return dpnp.moveaxis(a, 0, axis)
 
 
 def _scale_result(res, a_shape, norm, forward, index):
@@ -634,11 +617,12 @@ def dpnp_fft(a, forward, real, n=None, axis=-1, norm=None, out=None):
 
     if c2r:
         # input array should be Hermitian for c2r FFT
-        a = dpnp.moveaxis(a, axis, 0)
         a = _make_array_hermitian(
-            a, a.shape[0], dpnp.are_same_logical_tensors(a, a_orig)
+            a,
+            n=a.shape[0],
+            axis=axis,
+            copy_needed=dpnp.are_same_logical_tensors(a, a_orig),
         )
-        a = dpnp.moveaxis(a, 0, axis)
 
     return _fft(
         a,
@@ -687,11 +671,12 @@ def dpnp_fftn(a, forward, real, s=None, axes=None, norm=None, out=None):
     if len_axes == 1:
         a = _truncate_or_pad(a, (s[-1],), (axes[-1],))
         if c2r:
-            a = dpnp.moveaxis(a, axes[-1], 0)
             a = _make_array_hermitian(
-                a, a.shape[0], dpnp.are_same_logical_tensors(a, a_orig)
+                a,
+                n=a.shape[0],
+                axis=axes[-1],
+                copy_needed=dpnp.are_same_logical_tensors(a, a_orig),
             )
-            a = dpnp.moveaxis(a, 0, axes[-1])
         return _fft(
             a, norm, out, forward, in_place and c2c, c2c, axes[-1], a.ndim != 1
         )
@@ -743,11 +728,12 @@ def dpnp_fftn(a, forward, real, s=None, axes=None, norm=None, out=None):
         )
         a = _truncate_or_pad(a, (s[-1],), (axes[-1],))
         if c2r:
-            a = dpnp.moveaxis(a, axes[-1], 0)
             a = _make_array_hermitian(
-                a, a.shape[0], dpnp.are_same_logical_tensors(a, a_orig)
+                a,
+                n=a.shape[0],
+                axis=axes[-1],
+                copy_needed=dpnp.are_same_logical_tensors(a, a_orig),
             )
-            a = dpnp.moveaxis(a, 0, axes[-1])
         return _fft(
             a, norm, out, forward, in_place and c2c, c2c, axes[-1], a.ndim != 1
         )