diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 1d6ccb0..d3e5abe 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -29,11 +29,11 @@ jobs: - name: Install dependencies run: | python -m pip install --upgrade pip - python -m pip install .[dev] + python -m pip install -e .[dev] - name: Pytest run: | - pytest tests/ --cov + pytest tests/ --cov=pytorch_finufft lint: diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index d7babaf..5d41053 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -27,7 +27,7 @@ repos: - id: ruff - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.5.1 + rev: v1.8.0 hooks: - id: mypy exclude: ^tests/|^examples/ diff --git a/ci/Jenkinsfile b/ci/Jenkinsfile index 3f9b01a..87b8a49 100644 --- a/ci/Jenkinsfile +++ b/ci/Jenkinsfile @@ -16,53 +16,21 @@ pipeline { } environment { HOME = "$WORKSPACE" - LIBRARY_PATH = "$WORKSPACE/finufft/build" - LD_LIBRARY_PATH = "$WORKSPACE/finufft/build" } steps { - // TODO - reconsider install strategy once finufft/cufinufft 2.2 is released - checkout scmGit(branches: [[name: '*/master']], - extensions: [cloneOption(noTags: true, reference: '', shallow: true), - [$class: 'RelativeTargetDirectory', relativeTargetDir: 'finufft'], - cleanAfterCheckout()], - userRemoteConfigs: [[url: 'https://github.com/flatironinstitute/finufft']]) - - sh '''#!/bin/bash -ex - nvidia-smi - ''' - sh '''#!/bin/bash -ex - echo $HOME - ls - ''' - sh label: "build CUFINUFFT", script: '''#!/bin/bash -ex - cd finufft - # v100 cuda arch - cuda_arch="70" - mkdir build - cd build - cmake .. -DFINUFFT_USE_CUDA=ON \ - -DFINUFFT_USE_CPU=OFF \ - -DFINUFFT_BUILD_TESTS=OFF \ - -DCMAKE_CUDA_ARCHITECTURES="$cuda_arch" \ - -DBUILD_TESTING=OFF \ - -DCMAKE_BUILD_TYPE=RelWithDebInfo - cmake --build . -j 4 - ''' - sh 'python3 -m venv $HOME' sh label: "set up virtual environment", script: '''#!/bin/bash -ex source $HOME/bin/activate python3 -m pip install --upgrade pip # 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 -e .[dev] # can make [dev,cuda] once cufinufft released? + python3 -m pip install -e .[dev,cuda] ''' sh label: "run tests", script: '''#!/bin/bash -ex source $HOME/bin/activate - python3 -m pytest -k "cuda" tests/ --cov -v --error-for-skips + python3 -m pytest -k "cuda" tests/ --cov=pytorch_finufft -v --error-for-skips ''' } } diff --git a/pyproject.toml b/pyproject.toml index 5780f9f..bc2f6dc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -12,7 +12,7 @@ authors = [ {name = "Brian Ward", email="bward@flatironinstitute.org"}, ] license = { text = "MIT" } -dependencies = ["finufft>= 2.1", "torch >= 2", "numpy", "scipy"] +dependencies = ["finufft>= 2.2", "torch >= 2", "numpy", "scipy"] classifiers = [ "Development Status :: 3 - Alpha", "Intended Audience :: Developers", @@ -101,12 +101,3 @@ norecursedirs = [ ".git", "__pycache__", ] - -[tool.coverage.run] -source = ["tests"] - -[coverage.paths] -source = "pytorch_finufft" - -[coverage.run] -branch = true diff --git a/pytorch_finufft/functional.py b/pytorch_finufft/functional.py index dd46a4b..e4c1a6c 100644 --- a/pytorch_finufft/functional.py +++ b/pytorch_finufft/functional.py @@ -96,7 +96,7 @@ class FinufftType1(torch.autograd.Function): MODEORD_DEFAULT = 1 # note: FINUFFT default is 0 @staticmethod - def setup_context( # type: ignore[override] + def setup_context( ctx: Any, inputs: Tuple[ torch.Tensor, torch.Tensor, Any, Optional[Dict[str, Union[int, float]]] @@ -142,7 +142,7 @@ def forward( # type: ignore[override] batch_dims = values.shape[:-1] finufft_out = nufft_func( *points, - values.reshape(-1, values.shape[-1]).squeeze(), + values.reshape(-1, values.shape[-1]), output_shape, **finufftkwargs, ) @@ -237,10 +237,8 @@ def backward( # type: ignore[override] batched_values = values.reshape(nbatch, 1, values.shape[-1]) ramped_grad_output = ( - (coord_ramps * batched_grad_output * 1j * _i_sign) - .reshape(-1, *shape) - .squeeze() - ) + coord_ramps * batched_grad_output * 1j * _i_sign + ).reshape(-1, *shape) backprop_ramp = ( nufft_func(*points, ramped_grad_output, isign=_i_sign, **finufftkwargs) @@ -277,7 +275,7 @@ class FinufftType2(torch.autograd.Function): MODEORD_DEFAULT = 1 # note: FINUFFT default is 0 @staticmethod - def setup_context( # type: ignore[override] + def setup_context( ctx: Any, inputs: Tuple[ torch.Tensor, torch.Tensor, Optional[Dict[str, Union[int, float]]] @@ -324,7 +322,7 @@ def forward( # type: ignore[override] shape = targets.shape[-ndim:] finufft_out = nufft_func( *points, - targets.reshape(-1, *shape).squeeze(), + targets.reshape(-1, *shape), **finufftkwargs, ) finufft_out = finufft_out.reshape(*batch_dims, npoints) @@ -410,10 +408,8 @@ def backward( # type: ignore[override] coord_ramps = coordinate_ramps(shape, device) - ramped_targets = ( - (coord_ramps * batched_targets * 1j * _i_sign) - .reshape(-1, *shape) - .squeeze() # squeeze to work around finufft issue#367 + ramped_targets = (coord_ramps * batched_targets * 1j * _i_sign).reshape( + -1, *shape ) backprop_ramp = (