Provide PyTorch implementations by wrapping JAX functions

benchmarks/precompute_spherical.py

@@ -1,5 +1,6 @@
 """Benchmarks for precompute spherical transforms."""
 
+import jax
 import numpy as np
 from benchmarking import (
     BenchmarkSetup,
@@ -10,6 +11,7 @@
 
 import s2fft
 import s2fft.precompute_transforms
+from s2fft.utils import torch_wrapper
 
 L_VALUES = [8, 16, 32, 64, 128, 256]
 SPIN_VALUES = [0]
@@ -31,11 +33,17 @@ def setup_forward(method, L, sampling, spin, reality, recursion):
         sampling=sampling,
         reality=reality,
     )
-    kernel_function = (
-        s2fft.precompute_transforms.construct.spin_spherical_kernel_jax
-        if method == "jax"
-        else s2fft.precompute_transforms.construct.spin_spherical_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # arrays to a JAX arrays before converting to torch tensors which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flm = jax.numpy.asarray(flm)
+        f = jax.numpy.asarray(f)
+    if method.startswith("torch"):
+        flm, f = torch_wrapper.tree_map_jax_array_to_torch_tensor((flm, f))
+    kernel_function = s2fft.precompute_transforms.spherical._kernel_functions[method]
     kernel = kernel_function(
         L=L,
         spin=spin,
@@ -73,11 +81,16 @@ def setup_inverse(method, L, sampling, spin, reality, recursion):
         skip("Reality only valid for scalar fields (spin=0).")
     rng = np.random.default_rng()
     flm = s2fft.utils.signal_generator.generate_flm(rng, L, spin=spin, reality=reality)
-    kernel_function = (
-        s2fft.precompute_transforms.construct.spin_spherical_kernel_jax
-        if method == "jax"
-        else s2fft.precompute_transforms.construct.spin_spherical_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # array to a JAX array before converting to a torch tensor which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flm = jax.numpy.asarray(flm)
+    if method.startswith("torch"):
+        flm = torch_wrapper.jax_array_to_torch_tensor(flm)
+    kernel_function = s2fft.precompute_transforms.spherical._kernel_functions[method]
     kernel = kernel_function(
         L=L,
         spin=spin,

benchmarks/precompute_wigner.py

@@ -1,5 +1,6 @@
 """Benchmarks for precompute Wigner-d transforms."""
 
+import jax
 import numpy as np
 from benchmarking import (
     BenchmarkSetup,
@@ -10,6 +11,7 @@
 import s2fft
 import s2fft.precompute_transforms
 from s2fft.base_transforms import wigner as base_wigner
+from s2fft.utils import torch_wrapper
 
 L_VALUES = [16, 32, 64, 128, 256]
 N_VALUES = [2]
@@ -31,11 +33,17 @@ def setup_forward(method, L, N, L_lower, sampling, reality, mode):
         sampling=sampling,
         reality=reality,
     )
-    kernel_function = (
-        s2fft.precompute_transforms.construct.wigner_kernel_jax
-        if "jax" in method
-        else s2fft.precompute_transforms.construct.wigner_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # arrays to a JAX arrays before converting to torch tensors which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flmn = jax.numpy.asarray(flmn)
+        f = jax.numpy.asarray(f)
+    if method.startswith("torch"):
+        flmn, f = torch_wrapper.tree_map_jax_array_to_torch_tensor((flmn, f))
+    kernel_function = s2fft.precompute_transforms.wigner._kernel_functions[method]
     kernel = kernel_function(
         L=L, N=N, reality=reality, sampling=sampling, forward=True, mode=mode
     )
@@ -67,11 +75,16 @@ def forward(f, kernel, method, L, N, L_lower, sampling, reality, mode):
 def setup_inverse(method, L, N, L_lower, sampling, reality, mode):
     rng = np.random.default_rng()
     flmn = s2fft.utils.signal_generator.generate_flmn(rng, L, N, reality=reality)
-    kernel_function = (
-        s2fft.precompute_transforms.construct.wigner_kernel_jax
-        if method == "jax"
-        else s2fft.precompute_transforms.construct.wigner_kernel
-    )
+    # As torch method wraps JAX functions and converting NumPy array to torch tensor
+    # causes warning 'DLPack buffer is not aligned' about byte aligment on subsequently
+    # converting to JAX array using mutual DLPack support we first convert the NumPy
+    # arrays to a JAX arrays before converting to torch tensors which eliminates this
+    # warning
+    if method.startswith("jax") or method.startswith("torch"):
+        flmn = jax.numpy.asarray(flmn)
+    if method.startswith("torch"):
+        flmn = torch_wrapper.jax_array_to_torch_tensor(flmn)
+    kernel_function = s2fft.precompute_transforms.wigner._kernel_functions[method]
     kernel = kernel_function(
         L=L, N=N, reality=reality, sampling=sampling, forward=False, mode=mode
     )

notebooks/torch_frontend.ipynb

@@ -28,27 +28,21 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "This minimal tutorial demonstrates how to use the torch frontend for `S2FFT` to compute spherical harmonic transforms. Though `S2FFT` is primarily designed for JAX, this torch functionality is fully unit tested (including gradients) and can be used straightforwardly as a learnable layer within existing models."
+    "This minimal tutorial demonstrates how to use the torch frontend for `S2FFT` to compute spherical harmonic transforms. Though `S2FFT` is primarily designed for JAX, this torch functionality is fully unit tested (including gradients) and can be used straightforwardly as a learnable layer within existing models. As the torch functions wrap the JAX implementations we need to configure JAX to use 64-bit precision floating point types by default to ensure sufficient precision for the transforms - `S2FFT` will emit a warning if this has not been done."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": 2,
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "JAX is not using 64-bit precision. This will dramatically affect numerical precision at even moderate L.\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
+    "import jax\n",
+    "jax.config.update(\"jax_enable_x64\", True)\n",
     "import torch \n",
     "import numpy as np \n",
     "from s2fft.precompute_transforms.spherical import inverse, forward\n",
-    "from s2fft.precompute_transforms.construct import spin_spherical_kernel\n",
+    "from s2fft.precompute_transforms.construct import spin_spherical_kernel_torch\n",
     "from s2fft.utils import signal_generator"
    ]
   },
@@ -65,9 +59,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "L = 64                                                          # Spherical harmonic bandlimit\n",
-    "rng = np.random.default_rng(1234951510)                         # Random seed for signal generator\n",
-    "flm = signal_generator.generate_flm(rng, L, using_torch=True)   # Random set of spherical harmonic coefficients"
+    "L = 64 \n",
+    "rng = np.random.default_rng(1234951510)\n",
+    "flm = torch.from_numpy(signal_generator.generate_flm(rng, L))"
    ]
   },
   {
@@ -81,10 +75,18 @@
    "cell_type": "code",
    "execution_count": 4,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.\n"
+     ]
+    }
+   ],
    "source": [
-    "inverse_kernel = spin_spherical_kernel(L, using_torch=True, forward=False) \n",
-    "forward_kernel = spin_spherical_kernel(L, using_torch=True, forward=True) "
+    "inverse_kernel = spin_spherical_kernel_torch(L, forward=False) \n",
+    "forward_kernel = spin_spherical_kernel_torch(L, forward=True) "
    ]
   },
   {
@@ -135,7 +137,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Mean absolute error = 1.1866908936078849e-14\n"
+      "Mean absolute error = 2.8472981477378884e-14\n"
      ]
     }
    ],
@@ -146,7 +148,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3.10.4 ('s2fft')",
+   "display_name": "s2fft",
    "language": "python",
    "name": "python3"
   },
@@ -160,14 +162,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.0"
+   "version": "3.11.10"
   },
-  "orig_nbformat": 4,
-  "vscode": {
-   "interpreter": {
-    "hash": "3425e24474cbe920550266ea26b478634978cc419579f9dbcf479231067df6a3"
-   }
-  }
+  "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 2

pyproject.toml

@@ -30,7 +30,6 @@ dependencies = [
     "numpy>=1.20",
     "jax>=0.3.13",
     "jaxlib",
-    "torch",
 ]
 dynamic = [
     "version",
@@ -74,6 +73,10 @@ tests = [
     "pytest-cov",
     "so3",
     "pyssht",
+    "torch",
+]
+torch = [
+    "torch",
 ]
 
 [tool.scikit-build]

s2fft/precompute_transforms/construct.py

@@ -4,11 +4,10 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
-import torch
 
 from s2fft import recursions
 from s2fft.sampling import s2_samples as samples
-from s2fft.utils import quadrature, quadrature_jax
+from s2fft.utils import quadrature, quadrature_jax, torch_wrapper
 
 # Maximum spin number at which Price-McEwen recursion is sufficiently accurate.
 # For spins > PM_MAX_STABLE_SPIN one should default to the Risbo recursion.
@@ -22,7 +21,6 @@ def spin_spherical_kernel(
     sampling: str = "mw",
     nside: int = None,
     forward: bool = True,
-    using_torch: bool = False,
     recursion: str = "auto",
 ) -> np.ndarray:
     r"""
@@ -50,8 +48,6 @@ def spin_spherical_kernel(
         forward (bool, optional): Whether to provide forward or inverse shift.
             Defaults to False.
 
-        using_torch (bool, optional): Desired frontend functionality. Defaults to False.
-
         recursion (str, optional): Recursion to adopt. Supported recursion schemes include
             {"auto", "price-mcewen", "risbo"}. Defaults to "auto" which will detect the
             most appropriate recursion given the parameter configuration.
@@ -163,7 +159,7 @@ def spin_spherical_kernel(
             healpix_phase_shifts(L, nside, forward)[:, m_start_ind:],
         )
 
-    return torch.from_numpy(dl) if using_torch else dl
+    return dl
 
 
 def spin_spherical_kernel_jax(
@@ -329,6 +325,11 @@ def spin_spherical_kernel_jax(
     return dl
 
 
+spin_spherical_kernel_torch = torch_wrapper.wrap_as_torch_function(
+    spin_spherical_kernel_jax
+)
+
+
 def wigner_kernel(
     L: int,
     N: int,
@@ -337,7 +338,6 @@ def wigner_kernel(
     nside: int = None,
     forward: bool = False,
     mode: str = "auto",
-    using_torch: bool = False,
 ) -> np.ndarray:
     r"""
     Precompute the wigner-d kernel for Wigner transform.
@@ -368,8 +368,6 @@ def wigner_kernel(
             {"auto", "direct", "fft"}. Defaults to "auto" which will detect the
             most appropriate recursion given the parameter configuration.
 
-        using_torch (bool, optional): Desired frontend functionality. Defaults to False.
-
     Returns:
         np.ndarray: Transform kernel for Wigner transform.
 
@@ -468,7 +466,7 @@ def wigner_kernel(
             healpix_phase_shifts(L, nside, forward),
         )
 
-    return torch.from_numpy(dl) if using_torch else dl
+    return dl
 
 
 def wigner_kernel_jax(
@@ -611,6 +609,9 @@ def wigner_kernel_jax(
     return dl
 
 
+wigner_kernel_torch = torch_wrapper.wrap_as_torch_function(wigner_kernel_jax)
+
+
 def fourier_wigner_kernel(L: int) -> Tuple[np.ndarray, np.ndarray]:
     """
     Computes Fourier coefficients of the reduced Wigner d-functions and quadrature
@@ -667,6 +668,11 @@ def fourier_wigner_kernel_jax(L: int) -> Tuple[jnp.ndarray, jnp.ndarray]:
     return deltas, w
 
 
+fourier_wigner_kernel_torch = torch_wrapper.wrap_as_torch_function(
+    fourier_wigner_kernel_jax
+)
+
+
 def healpix_phase_shifts(L: int, nside: int, forward: bool = False) -> np.ndarray:
     r"""
     Generates a phase shift vector for HEALPix for all :math:`\theta` rings.