format

Author: ASKabalan

lib/src/extensions.cc

@@ -251,7 +251,7 @@ ffi::Error healpix_backward(cudaStream_t stream, ffi::Buffer<T> input, ffi::Resu
  */
 template <ffi::DataType T>
 s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, bool reality, bool forward,
-                                 bool normalize, bool adjoint, bool must_exist , size_t& work_size) {
+                                 bool normalize, bool adjoint, bool must_exist, size_t& work_size) {
     using fft_complex_type = fft_complex_t<T>;
     // Step 1: Determine FFT normalization type based on forward/normalize flags.
     s2fftKernels::fft_norm norm = s2fftKernels::fft_norm::NONE;
@@ -285,7 +285,7 @@ s2fftDescriptor build_descriptor(int64_t nside, int64_t harmonic_band_limit, boo
     if (hr == S_OK) {
         executor->Initialize(descriptor);
     }
-    // Make sure workspace is set 
+    // Make sure workspace is set
     assert(executor->m_work_size > 0 && "S2FFT INTERNAL ERROR: Workspace size is zero after initialization.");
     work_size = executor->m_work_size;
     // Step 7: Return the created descriptor.
@@ -320,8 +320,8 @@ ffi::Error healpix_fft_cuda(cudaStream_t stream, int64_t nside, int64_t harmonic
                             ffi::Result<ffi::Buffer<ffi::DataType::S64>> callback_params) {
     // Step 1: Build the s2fftDescriptor based on the input parameters.
     size_t work_size = 0;  // Variable to hold the workspace size
-    s2fftDescriptor descriptor =
-            build_descriptor<T>(nside, harmonic_band_limit, reality, forward, normalize, adjoint, true , work_size);
+    s2fftDescriptor descriptor = build_descriptor<T>(nside, harmonic_band_limit, reality, forward, normalize,
+                                                     adjoint, true, work_size);
 
     // Step 2: Dispatch to either forward or backward transform based on the 'forward' flag.
     if (forward) {

s2fft/utils/healpix_ffts.py

@@ -587,6 +587,7 @@ def _healpix_fft_cuda_abstract(f, L, nside, reality, fft_type, norm, adjoint):
 
     Returns:
         Tuple of ShapedArray objects for output, workspace, and callback parameters.
+
     """
     # Step 1: Get lowering information (double precision, forward/backward, normalize).
     is_double, forward, normalize = _get_lowering_info(fft_type, norm, f.dtype)
@@ -632,7 +633,7 @@ def _healpix_fft_cuda_abstract(f, L, nside, reality, fft_type, norm, adjoint):
         shape=batch_shape + workspace_shape, dtype=workspace_dtype
     )
     params_eval = ShapedArray(shape=batch_shape + params_shape, dtype=np.int64)
-    
+
     # Step 6: Return the ShapedArray objects.
     return (
         f.update(shape=out_shape, dtype=f.dtype),
@@ -666,6 +667,7 @@ def _healpix_fft_cuda_lowering(ctx, f, *, L, nside, reality, fft_type, norm, adj
 
     Returns:
         The result of the FFI call.
+
     """
     # Step 1: Check if CUDA support is compiled in.
     if not _s2fft.COMPILED_WITH_CUDA:
@@ -715,6 +717,7 @@ def _healpix_fft_cuda_batching_rule(
 
     Returns:
         Tuple of (output, output_batch_axes).
+
     """
     # Step 1: Unpack batched arguments and batching axes.
     (x,) = batched_args
@@ -772,6 +775,7 @@ def _healpix_fft_cuda_transpose(
 
     Returns:
         The adjoint of the input.
+
     """
     # Step 1: Invert the FFT type and normalization for the adjoint operation.
     fft_type = "backward" if fft_type == "forward" else "forward"
@@ -901,4 +905,4 @@ def healpix_ifft_cuda(
     "jax": healpix_ifft_jax,
     "cuda": healpix_ifft_cuda,
     "torch": healpix_ifft_torch,
-}
+}