CI: tweaks

flatironinstitute · Oct 6, 2023 · f431920 · f431920
1 parent 03c79d6
commit f431920
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 18 deletions.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -16,7 +16,7 @@ pipeline {
       }
       environment {
     HOME = "$WORKSPACE"
-    PYBIN = "/opt/python/cp38-cp38/bin"
+    PYBIN = "/opt/python/cp39-cp39/bin"
     LIBRARY_PATH = "$WORKSPACE/finufft/build"
     LD_LIBRARY_PATH = "$WORKSPACE/finufft/build"
       }
@@ -42,6 +42,7 @@ pipeline {
         cuda_arch="70"
 
         cmake -B build . -DFINUFFT_USE_CUDA=ON \
+                         -DFINUFFT_USE_CPU=OFF \
                          -DFINUFFT_BUILD_TESTS=OFF \
                          -DCMAKE_CUDA_ARCHITECTURES="$cuda_arch" \
                          -DBUILD_TESTING=ON
@@ -56,11 +57,10 @@ pipeline {
       # we could also move pytorch install inside docker
       python3 -m pip install "torch~=2.1.0" --index-url https://download.pytorch.org/whl/cu118
       python3 -m pip install finufft/python/cufinufft
-      python3 -m pip install finufft/python/finufft
 
       python3 -m pip install -e .[dev]
 
-      python3 -m pytest -k "cuda" tests/ --cov
+      python3 -m pytest -k "cuda" tests/ --cov -v
     '''
       }
     }

diff --git a/pytorch_finufft/functional.py b/pytorch_finufft/functional.py
@@ -4,16 +4,28 @@
 
 from typing import Any, Dict, Optional, Tuple, Union
 
-import finufft
 import numpy as np
+import torch
+
+try:
+    import finufft
+
+    FINUFFT_AVAIL = True
+except ImportError:
+    FINUFFT_AVAIL = False
 
 try:
     import cufinufft
 
     CUFINUFFT_AVAIL = True
 except ImportError:
     CUFINUFFT_AVAIL = False
-import torch
+
+if not (FINUFFT_AVAIL or CUFINUFFT_AVAIL):
+    raise ImportError(
+        "No FINUFFT implementation available. "
+        "Install either finufft or cufinufft and ensure they are importable."
+    )
 
 import pytorch_finufft._err as err
 
@@ -1631,7 +1643,7 @@ def forward(
         ctx: Any,
         points: torch.Tensor,
         values: torch.Tensor,
-        output_shape: Union[int, tuple[int, int], tuple[int, int, int]],
+        output_shape: Union[int, Tuple[int, int], Tuple[int, int, int]],
         out: Optional[torch.Tensor] = None,
         fftshift: bool = False,
         finufftkwargs: dict[str, Union[int, float]] = None,
@@ -1693,21 +1705,21 @@ def forward(
     @staticmethod
     def backward(
         ctx: Any, grad_output: torch.Tensor
-    ) -> tuple[Union[torch.Tensor, None], ...]:
+    ) -> Tuple[Union[torch.Tensor, None], ...]:
         """
-        Implements derivatives wrt. each argument in the forward method.
+         Implements derivatives wrt. each argument in the forward method.
 
-        Parameters
-        ----------
-        ctx : Any
-            Pytorch context object.
-        grad_output : torch.Tensor
-            Backpass gradient output
+         Parameters
+         ----------
+         ctx : Any
+             Pytorch context object.
+         grad_output : torch.Tensor
+             Backpass gradient output
 
-        Returns
-        -------
-        tuple[Union[torch.Tensor, None], ...]
-            A tuple of derivatives wrt. each argument in the forward method
+         Returns
+         -------
+        Tuple[Union[torch.Tensor, None], ...]
+             A tuple of derivatives wrt. each argument in the forward method
         """
         _i_sign = -1 * ctx.isign
         _mode_ordering = ctx.mode_ordering