Merge pull request #231 from astro-informatics/mmg/cuda-primitive

Updates to custom CUDA HEALPix FFT primitive

Author: ASKabalan

.github/workflows/docs.yml

@@ -26,10 +26,7 @@ jobs:
       run: |
         sudo apt install pandoc
         python -m pip install --upgrade pip
-        pip install jaxlib
-        pip install -r requirements/requirements-core.txt
-        pip install -r requirements/requirements-docs.txt
-        pip install .
+        pip install .[docs]
   
     - name: Build Documentation
       run: |

.github/workflows/tests.yml

@@ -28,9 +28,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -r requirements/requirements-tests.txt
-        pip install -r requirements/requirements-core.txt
-        pip install .
+        pip install .[tests]
     
     - name: Run tests
       run: |

MANIFEST.in

@@ -1,6 +1,5 @@
 graft s2fft
 graft tests
-graft requirements
 graft docs
 
 include .pip_readme.rst
@@ -14,4 +13,4 @@ exclude *.log
 recursive-exclude docs/_build *
 recursive-exclude notebooks *
 
-global-exclude *.py[cod] __pycache__ *.so *.dylib .ipynb_checkpoints/** *~
+global-exclude *.py[cod] __pycache__ *.so *.dylib .ipynb_checkpoints/** *~

README.md

@@ -290,4 +290,13 @@ it will be of use to a wider community.
 Copyright 2023 Matthew Price, Jason McEwen and contributors.
 
 `S2FFT` is free software made available under the MIT License. For
-details see the LICENSE file.
+details see the [`LICENCE.txt`](LICENCE.txt) file.
+
+The file [`lib/include/kernel_helpers.h`](lib/include/kernel_helpers.h) is adapted from
+[code](https://github.com/dfm/extending-jax/blob/c33869665236877a2ae281f3f5dbff579e8f5b00/lib/kernel_helpers.h) in [a tutorial on extending JAX](https://github.com/dfm/extending-jax) by 
+[Dan Foreman-Mackey](https://github.com/dfm) and licensed under a [MIT license](https://github.com/dfm/extending-jax/blob/371dca93c6405368fa8e71690afd3968d75f4bac/LICENSE).
+
+The file [`lib/include/kernel_nanobind_helpers.h`](lib/include/kernel_nanobind_helpers.h)
+is adapted from [code](https://github.com/jax-ml/jax/blob/3d389a7fb440c412d95a1f70ffb91d58408247d0/jaxlib/kernel_nanobind_helpers.h) 
+by the [JAX](https://github.com/jax-ml/jax) authors 
+and licensed under a [Apache-2.0 license](https://github.com/jax-ml/jax/blob/3d389a7fb440c412d95a1f70ffb91d58408247d0/LICENSE). 

lib/include/kernel_helpers.h

@@ -1,3 +1,30 @@
+// Adapted from code in a tutorial by Dan Foreman-Mackey
+// https://github.com/dfm/extending-jax/blob/c33869665236877a2ae281/lib/kernel_helpers.h
+//
+// Original license:
+//
+// MIT License
+//
+// Copyright (c) 2021 Dan Foreman-Mackey
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in all
+// copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+// SOFTWARE.
+
 // This header is not specific to our application and you'll probably want
 // something like this for any extension you're building. This includes the
 // infrastructure needed to serialize descriptors that are used with the

lib/include/kernel_nanobind_helpers.h

@@ -1,3 +1,6 @@
+// Adapted from code by JAX authors
+// https://github.com/jax-ml/jax/blob/3d389a7fb440c412d/jaxlib/kernel_nanobind_helpers.h
+
 /* Copyright 2019 The JAX Authors.
 
 Licensed under the Apache License, Version 2.0 (the "License");