Tidy up the backwards pass of type 1

pytorch_finufft/functional.py

@@ -776,7 +776,7 @@ def forward(
         return finufft_out
 
     @staticmethod
-    def backward(
+    def backward(  # type: ignore[override]
         ctx: Any, grad_output: torch.Tensor
     ) -> Tuple[Union[torch.Tensor, None], ...]:
         """
@@ -799,55 +799,42 @@ def backward(
         finufftkwargs = ctx.finufftkwargs
 
         points, values = ctx.saved_tensors
-
-        start_points = -(np.array(grad_output.shape) // 2)
-        end_points = start_points + grad_output.shape
-        slices = tuple(
-            slice(start, end) for start, end in zip(start_points, end_points)
-        )
-
-        # CPU idiosyncracy that needs to be done differently
-        coord_ramps = torch.from_numpy(np.mgrid[slices]).to(points.device)
+        device = points.device
 
         grads_points = None
         grad_values = None
 
         ndim = points.shape[0]
 
-        nufft_func = get_nufft_func(ndim, 2, points.device.type)
+        nufft_func = get_nufft_func(ndim, 2, device.type)
+
+        if any(ctx.needs_input_grad) and _mode_ordering:
+            grad_output = torch.fft.fftshift(grad_output)
 
         if ctx.needs_input_grad[0]:
             # wrt points
-
-            if _mode_ordering:
-                coord_ramps = torch.fft.ifftshift(
-                    coord_ramps, dim=tuple(range(1, ndim + 1))
+            start_points = -(torch.tensor(grad_output.shape, device=device) // 2)
+            end_points = start_points + torch.tensor(grad_output.shape, device=device)
+            coord_ramps = torch.stack(
+                torch.meshgrid(
+                    *(
+                        torch.arange(start, end, device=device)
+                        for start, end in zip(start_points, end_points)
+                    ),
+                    indexing="ij",
                 )
+            )
 
-            ramped_grad_output = coord_ramps * grad_output[np.newaxis] * 1j * _i_sign
-
-            grads_points = []
-            for ramp in ramped_grad_output:  # we can batch this into finufft
-                if _mode_ordering:
-                    ramp = torch.fft.fftshift(ramp)
-
-                backprop_ramp = nufft_func(
-                    *points,
-                    ramp,
-                    isign=_i_sign,
-                    **finufftkwargs,
-                )
-
-                grad_points = (backprop_ramp.conj() * values).real
-
-                grads_points.append(grad_points)
-
-            grads_points = torch.stack(grads_points)
+            # we can't batch in 1d case so we squeeze and fix up the ouput later
+            ramped_grad_output = (
+                coord_ramps * grad_output[np.newaxis] * 1j * _i_sign
+            ).squeeze()
+            backprop_ramp = nufft_func(
+                *points, ramped_grad_output, isign=_i_sign, **finufftkwargs
+            )
+            grads_points = torch.atleast_2d((backprop_ramp.conj() * values).real)
 
         if ctx.needs_input_grad[1]:
-            if _mode_ordering:
-                grad_output = torch.fft.fftshift(grad_output)
-
             grad_values = nufft_func(
                 *points,
                 grad_output,