torch-bessel 0.0.9__tar.gz

torch_bessel-0.0.9/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Ho Yin Chau
+
+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.

torch_bessel-0.0.9/PKG-INFO

@@ -0,0 +1,50 @@
+Metadata-Version: 2.4
+Name: torch_bessel
+Version: 0.0.9
+Summary: PyTorch extension package for Bessel functions with arbitrary real order and complex inputs
+Home-page: https://github.com/hchau630/torch-bessel
+Author: Ho Yin Chau
+Requires-Python: >= 3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# About
+PyTorch extension package for modified Bessel functions of the second kind with complex inputs
+
+# Install
+Currently only supports Linux (with CUDA 12.4) or MacOS (Apple silicon, cpu only) with python >= 3.9, <= 3.12.
+```
+pip install torch-bessel
+```
+
+UPDATE (May 26, 2026): It seems the package might no longer be compatible with the latest pytorch version (>= 2.12.0), so try installing an earlier version of pytorch if you encounter issues.
+
+UPDATE (May 28, 2026): On Apple silicon, importing this package with pytorch version >= 2.9.0 causes a crash (bus error: 10) when exiting a python program, though the program runs normally prior to that. So consider installing torch < 2.9.0 on Apple silicon.
+
+# Example
+```
+import torch_bessel
+
+real, imag = torch.randn(2, 5, device="cuda")
+z = torch.complex(real.abs(), imag)  # correctness for inputs in the left-half complex plane is not gauranteed.
+torch_bessel.ops.modified_bessel_k0(z)
+```
+
+# Implemented functions
+- `modified_bessel_k0`: Same as `torch.special.modified_bessel_k0`, but also handles backpropagation and complex inputs on cpu and cuda. Correctness is guaranteed on the right-half complex plane. On cuda, `torch.chalf` inputs are also supported, though the underlying cuda kernel just upcasts `chalf` to `cfloat` (note that this uses no extra GPU memory, as opposed to manually casting torch.chalf to torch.cfloat before calling `modified_bessel_k0` which doubles the GPU memory used). On the left-half complex plane, function output appears mostly correct, but with small numerical errors for certain inputs. On the negative real line, output is NaN.
+- `modified_bessel_k1`: Same as `torch.special.modified_bessel_k1`, but also handles complex inputs on cpu and cuda. Backpropagation not implemented, but this can be easily manually implemented yourself by writing a torch.autograd.Function using the recurrence properties of bessel functions. Same caveats as `modified_bessel_k0` apply.
+
+# WIP
+- `modified_bessel_kv`: Analogue of `scipy.special.kv`.
+
+# Benchmarks
+Benchmarking performed with the `asv` package. Results can be viewed at https://hchau630.github.io/torch-bessel.

torch_bessel-0.0.9/README.md

@@ -0,0 +1,31 @@
+# About
+PyTorch extension package for modified Bessel functions of the second kind with complex inputs
+
+# Install
+Currently only supports Linux (with CUDA 12.4) or MacOS (Apple silicon, cpu only) with python >= 3.9, <= 3.12.
+```
+pip install torch-bessel
+```
+
+UPDATE (May 26, 2026): It seems the package might no longer be compatible with the latest pytorch version (>= 2.12.0), so try installing an earlier version of pytorch if you encounter issues.
+
+UPDATE (May 28, 2026): On Apple silicon, importing this package with pytorch version >= 2.9.0 causes a crash (bus error: 10) when exiting a python program, though the program runs normally prior to that. So consider installing torch < 2.9.0 on Apple silicon.
+
+# Example
+```
+import torch_bessel
+
+real, imag = torch.randn(2, 5, device="cuda")
+z = torch.complex(real.abs(), imag)  # correctness for inputs in the left-half complex plane is not gauranteed.
+torch_bessel.ops.modified_bessel_k0(z)
+```
+
+# Implemented functions
+- `modified_bessel_k0`: Same as `torch.special.modified_bessel_k0`, but also handles backpropagation and complex inputs on cpu and cuda. Correctness is guaranteed on the right-half complex plane. On cuda, `torch.chalf` inputs are also supported, though the underlying cuda kernel just upcasts `chalf` to `cfloat` (note that this uses no extra GPU memory, as opposed to manually casting torch.chalf to torch.cfloat before calling `modified_bessel_k0` which doubles the GPU memory used). On the left-half complex plane, function output appears mostly correct, but with small numerical errors for certain inputs. On the negative real line, output is NaN.
+- `modified_bessel_k1`: Same as `torch.special.modified_bessel_k1`, but also handles complex inputs on cpu and cuda. Backpropagation not implemented, but this can be easily manually implemented yourself by writing a torch.autograd.Function using the recurrence properties of bessel functions. Same caveats as `modified_bessel_k0` apply.
+
+# WIP
+- `modified_bessel_kv`: Analogue of `scipy.special.kv`.
+
+# Benchmarks
+Benchmarking performed with the `asv` package. Results can be viewed at https://hchau630.github.io/torch-bessel.

torch_bessel-0.0.9/benchmarks/benchmarks.py

@@ -0,0 +1,143 @@
+import torch
+from asv_runner.benchmarks.mark import skip_benchmark_if
+
+import torch_bessel
+
+
+def _setup(
+    n, is_real, singularity, dtype=torch.float, requires_grad=False, device="cpu"
+):
+    kwargs = {"dtype": dtype, "requires_grad": requires_grad, "device": device}
+    real = torch.randn(n, **kwargs).abs()
+    if is_real:
+        args = (real, singularity)
+    else:
+        imag = torch.randn(n, **kwargs)
+        args = (torch.complex(real, imag), singularity)
+    return args
+
+
+def _setup_k1(n, is_real, dtype=torch.float, device="cpu"):
+    kwargs = {"dtype": dtype, "device": device}
+    real = torch.randn(n, **kwargs).abs()
+    if is_real:
+        args = (real,)
+    else:
+        imag = torch.randn(n, **kwargs)
+        args = (torch.complex(real, imag),)
+    return args
+
+
+class ModifiedBesselK0ForwardCPU:
+    params = (
+        [10_000, 100_000, 1_000_000],
+        [False, True],
+        [None, 0.0],
+        [torch.float32, torch.float64],
+        [False, True],
+    )
+    param_names = ["n", "is_real", "singularity", "dtype", "requires_grad"]
+
+    def setup(self, n, is_real, singularity, dtype, requires_grad):
+        self.args = _setup(n, is_real, singularity, dtype, requires_grad)
+
+    def time_modified_bessel_k0_forward_cpu(
+        self, n, is_real, singularity, dtype, requires_grad
+    ):
+        torch_bessel.ops.modified_bessel_k0(*self.args)
+
+
+class ModifiedBesselK0ForwardCUDA:
+    params = (
+        [10_000, 100_000, 1_000_000],
+        [False, True],
+        [None, 0.0],
+        [torch.float32, torch.float64],
+        [False, True],
+    )
+    param_names = ["n", "is_real", "singularity", "dtype", "requires_grad"]
+
+    def setup(self, n, is_real, singularity, dtype, requires_grad):
+        self.args = _setup(n, is_real, singularity, dtype, requires_grad, device="cuda")
+
+    @skip_benchmark_if(not torch.cuda.is_available())
+    def time_modified_bessel_k0_forward_cuda(
+        self, n, is_real, singularity, dtype, requires_grad
+    ):
+        torch.cuda.synchronize()
+        torch_bessel.ops.modified_bessel_k0(*self.args)
+        torch.cuda.synchronize()
+
+
+class ModifiedBesselK0BackwardCPU:
+    warmup_time = 0.0  # for some reason backward is called multiple times if not 0
+    number = 1  # Avoids calling backward multiple times
+    params = (
+        [10_000, 100_000, 1_000_000],
+        [False, True],
+        [None, 0.0],
+        [torch.float32, torch.float64],
+    )
+    param_names = ["n", "is_real", "singularity", "dtype"]
+
+    def setup(self, n, is_real, singularity, dtype):
+        args = _setup(n, is_real, singularity, dtype, requires_grad=True)
+        self.out = torch_bessel.ops.modified_bessel_k0(*args).norm()
+
+    def time_modified_bessel_k0_backward_cpu(self, n, is_real, singularity, dtype):
+        self.out.backward()
+
+
+class ModifiedBesselK0BackwardCUDA:
+    warmup_time = 0.0  # for some reason backward is called multiple times if not 0
+    number = 1  # Avoids calling backward multiple times
+    params = (
+        [10_000, 100_000, 1_000_000],
+        [False, True],
+        [None, 0.0],
+        [torch.float32, torch.float64],
+    )
+    param_names = ["n", "is_real", "singularity", "dtype"]
+
+    def setup(self, n, is_real, singularity, dtype):
+        args = _setup(n, is_real, singularity, dtype, requires_grad=True, device="cuda")
+        self.out = torch_bessel.ops.modified_bessel_k0(*args).norm()
+
+    @skip_benchmark_if(not torch.cuda.is_available())
+    def time_modified_bessel_k0_backward_cuda(self, n, is_real, singularity, dtype):
+        torch.cuda.synchronize()
+        self.out.backward()
+        torch.cuda.synchronize()
+
+
+class ModifiedBesselK1ForwardCPU:
+    params = (
+        [10_000, 100_000, 1_000_000],
+        [False, True],
+        [torch.float32, torch.float64],
+    )
+    param_names = ["n", "is_real", "dtype"]
+
+    def setup(self, n, is_real, dtype):
+        self.args = _setup_k1(n, is_real, dtype)
+
+    def time_modified_bessel_k1_forward_cpu(self, n, is_real, dtype):
+        torch_bessel.ops.modified_bessel_k1(*self.args)
+
+
+class ModifiedBesselK1ForwardCUDA:
+    params = (
+        [10_000, 100_000, 1_000_000],
+        [False, True],
+        [torch.float32, torch.float64],
+    )
+    param_names = ["n", "is_real", "dtype"]
+
+    def setup(self, n, is_real, dtype):
+        self.args = _setup_k1(n, is_real, dtype, device="cuda")
+
+    @skip_benchmark_if(not torch.cuda.is_available())
+    def time_modified_bessel_k1_forward_cuda(self, n, is_real, dtype):
+        torch.cuda.synchronize()
+        torch_bessel.ops.modified_bessel_k1(*self.args)
+        torch.cuda.synchronize()

torch_bessel-0.0.9/pyproject.toml

@@ -0,0 +1,9 @@
+[build-system]
+requires = [
+    "setuptools",
+    "torch>=2.4.0",
+    # "torch>=2.7.0; platform_system == 'Darwin'",
+    # "torch>=2.4.0; platform_system != 'Darwin'",
+    "ninja",
+]
+build-backend = "setuptools.build_meta"

torch_bessel-0.0.9/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

torch_bessel-0.0.9/setup.py

@@ -0,0 +1,80 @@
+import os
+import glob
+
+from setuptools import find_packages, setup
+
+from torch.utils.cpp_extension import (
+    CppExtension,
+    CUDAExtension,
+    BuildExtension,
+    CUDA_HOME,
+)
+
+library_name = "torch_bessel"
+
+
+def get_extensions():
+    debug_mode = os.getenv("DEBUG", "0") == "1"
+    use_cuda = os.getenv("USE_CUDA", "1") == "1"
+    if debug_mode:
+        print("Compiling in debug mode")
+
+    use_cuda = use_cuda and CUDA_HOME is not None
+    extension = CUDAExtension if use_cuda else CppExtension
+
+    extra_link_args = []
+    extra_compile_args = {
+        "cxx": [
+            "-O3" if not debug_mode else "-O0",
+            "-std=c++17",
+            "-fdiagnostics-color=always",
+            "-DPy_LIMITED_API=0x030A0000",
+        ],
+        "nvcc": [
+            "-O3" if not debug_mode else "-O0",
+            "--extended-lambda",
+        ],
+    }
+    if debug_mode:
+        extra_compile_args["cxx"].append("-g")
+        extra_compile_args["nvcc"].append("-g")
+        extra_link_args.extend(["-O0", "-g"])
+
+    this_dir = os.path.dirname(os.path.curdir)
+    extensions_dir = os.path.join(this_dir, library_name, "csrc")
+    sources = list(glob.glob(os.path.join(extensions_dir, "*.cpp")))
+
+    extensions_cuda_dir = os.path.join(extensions_dir, "cuda")
+    cuda_sources = list(glob.glob(os.path.join(extensions_cuda_dir, "*.cu")))
+
+    if use_cuda:
+        sources += cuda_sources
+
+    ext_modules = [
+        extension(
+            f"{library_name}._C",
+            sources,
+            extra_compile_args=extra_compile_args,
+            extra_link_args=extra_link_args,
+            py_limited_api=True,
+        )
+    ]
+
+    return ext_modules
+
+
+setup(
+    name=library_name,
+    version="0.0.9",
+    author="Ho Yin Chau",
+    packages=find_packages(),
+    ext_modules=get_extensions(),
+    install_requires=["torch"],
+    python_requires=">= 3.10",
+    description="PyTorch extension package for Bessel functions with arbitrary real order and complex inputs",
+    long_description=open("README.md").read(),
+    long_description_content_type="text/markdown",
+    url="https://github.com/hchau630/torch-bessel",
+    cmdclass={"build_ext": BuildExtension},
+    options={"bdist_wheel": {"py_limited_api": "cp310"}},
+)

torch_bessel-0.0.9/tests/test_extension.py

@@ -0,0 +1,281 @@
+import torch
+from torch.testing._internal.common_utils import TestCase
+from torch.testing._internal.optests import opcheck
+import pytest
+from scipy import special
+
+import torch_bessel
+
+
+def reference_modified_bessel_k0(z, singularity=None):
+    device = z.device
+    dtype = z.dtype
+    if dtype is torch.chalf:
+        z = z.to(torch.cfloat)
+    out = special.kv(0.0, z.detach().cpu()).to(device).to(dtype)
+    if singularity is not None:
+        out = out.where(z != 0, singularity)
+    return out
+
+
+def reference_modified_bessel_k1(z):
+    device = z.device
+    dtype = z.dtype
+    if dtype is torch.chalf:
+        z = z.to(torch.cfloat)
+    out = special.kv(1.0, z.detach().cpu()).to(device).to(dtype)
+    return out
+
+
+class TestBesselK0(TestCase):
+    def sample_inputs(self, device, *, requires_grad=False):
+        def make_z(*size, dtype=None):
+            kwargs = dict(dtype=dtype, device=device, requires_grad=requires_grad)
+            real = torch.randn(size, **kwargs).abs()
+            imag = torch.randn(size, **kwargs)
+            return torch.complex(real, imag)
+
+        out = [
+            [make_z(50, dtype=torch.double)],
+            [make_z(50, dtype=torch.double).real],
+            [make_z(50)],
+            [make_z(50).real],
+            [torch.tensor(70.1162 + 89.0190j)],  # used to fail on this input
+        ]
+        if device == "cuda":
+            # half precision is only supported on CUDA
+            out.append([make_z(50, dtype=torch.half)])
+
+        return out
+
+    def grid_inputs(self, device, *, requires_grad=False):
+        def make_z(*args, dtype=torch.float):
+            real = torch.tensor([0, *torch.logspace(*args, dtype=torch.double)])
+            imag = torch.logspace(*args, dtype=torch.double)
+
+            # Don't test on subnormal numbers, since the AMOS code doesn't consider them
+            real[real < torch.finfo(dtype).tiny] = 0.0
+            imag[imag < torch.finfo(dtype).tiny] = 0.0
+
+            kwargs = dict(dtype=dtype, device=device, requires_grad=requires_grad)
+            real = torch.tensor(real, **kwargs)
+            imag = torch.tensor([*(-imag.flip(0)), 0, *imag], **kwargs)
+            real, imag = real[:, None], imag[None, :]
+            return torch.complex(real, imag)
+
+        out = [
+            [make_z(-350, 350, 75, dtype=torch.double)],
+            [make_z(-50, 50, 75)],
+            [make_z(-5, 5, 75, dtype=torch.double)],  # scipy is buggy on this input
+            [make_z(-5, 5, 75)],  # scipy is buggy on this input
+            [make_z(-50, 50, 75), 1.0],
+            [make_z(-50, 50, 75), torch.randn((76, 151), device=device)],
+        ]
+        if device == "cuda":
+            # half precision is only supported on CUDA
+            out.append([make_z(-5, 5, 75, dtype=torch.half)])
+        return out
+
+    def _test_correctness(self, device):
+        samples = (
+            self.sample_inputs(device)
+            + self.grid_inputs(device)
+            + self.sample_inputs(device, requires_grad=True)
+            + self.grid_inputs(device, requires_grad=True)
+        )
+        for args in samples:
+            result = torch_bessel.ops.modified_bessel_k0(*args)
+            expected = reference_modified_bessel_k0(*args)
+            if expected.dtype in {torch.float, torch.cfloat, torch.chalf}:
+                # ierr = 4, complete loss of significance
+                expected[args[0].abs() > 4194303.98419452] = torch.nan
+                # ierr = 2, overflow
+                mask = (args[0].abs() < 1.1754944e-35) & (args[0] != 0)
+                if expected.is_complex():
+                    expected[mask] = torch.nan
+                    expected[mask & (args[0].imag == 0) & (args[0] != 0)] = torch.inf
+                else:
+                    expected[mask & (args[0] != 0)] = torch.inf
+            # fix scipy bug. See https://github.com/scipy/xsf/issues/46
+            if args[0].is_complex():
+                mask = (
+                    (args[0].real > 685)
+                    & (args[0].real < 690)
+                    & (args[0].imag.abs() > 685)
+                    & (args[0].imag.abs() < 1.1e5)
+                )
+                expected[mask] = 0.0  # scipy returns incorrect/NaN output
+
+            torch.testing.assert_close(result, expected, equal_nan=True)
+
+    def test_correctness_cpu(self):
+        self._test_correctness("cpu")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
+    def test_correctness_cuda(self):
+        self._test_correctness("cuda")
+
+    def _test_gradients(self, device):
+        samples = self.sample_inputs(device, requires_grad=True)
+        for args in samples:
+            if (
+                args[0].dtype in {torch.double, torch.complex128}
+                and args[0].requires_grad
+            ):
+                torch.autograd.gradcheck(torch_bessel.ops.modified_bessel_k0, args)
+
+    def test_gradients_cpu(self):
+        self._test_gradients("cpu")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
+    def test_gradients_cuda(self):
+        self._test_gradients("cuda")
+
+    def _opcheck(self, device):
+        # Use opcheck to check for incorrect usage of operator registration APIs
+        samples = self.sample_inputs(device, requires_grad=False)
+        samples.extend(self.sample_inputs(device, requires_grad=True))
+        for args in samples:
+            if not args[0].is_complex():
+                continue
+            opcheck(
+                torch.ops.torch_bessel.modified_bessel_k0_complex_forward_backward.default,
+                args,
+            )
+
+    def test_opcheck_cpu(self):
+        self._opcheck("cpu")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
+    def test_opcheck_cuda(self):
+        self._opcheck("cuda")
+
+
+@pytest.mark.parametrize(
+    "singularity, in_dims",
+    [
+        (None, (0, None)),
+        (0.0, (0, None)),
+        ((4, 3), (0, 1)),
+    ],
+)
+def test_vmap(singularity, in_dims):
+    z = torch.randn(3, 4).abs() + torch.randn(3, 4) * 1j
+    if isinstance(singularity, tuple):
+        singularity = torch.randn(singularity)
+    func = torch.func.vmap(torch_bessel.ops.modified_bessel_k0, in_dims=in_dims)
+    out = func(z, singularity)
+    if not isinstance(singularity, torch.Tensor):
+        expected = torch_bessel.ops.modified_bessel_k0(z, singularity)
+    else:
+        expected = torch_bessel.ops.modified_bessel_k0(z, singularity.t())
+    torch.testing.assert_close(out, expected)
+
+
+class TestBesselK1(TestCase):
+    def sample_inputs(self, device, *, requires_grad=False):
+        def make_z(*size, dtype=None):
+            kwargs = dict(dtype=dtype, device=device, requires_grad=requires_grad)
+            real = torch.randn(size, **kwargs).abs()
+            imag = torch.randn(size, **kwargs)
+            return torch.complex(real, imag)
+
+        out = [
+            [make_z(50, dtype=torch.double)],
+            [make_z(50, dtype=torch.double).real],
+            [make_z(50)],
+            [make_z(50).real],
+            [torch.tensor(70.1162 + 89.0190j)],  # used to fail on this input
+        ]
+        if device == "cuda":
+            # half precision is only supported on CUDA
+            out.append([make_z(50, dtype=torch.half)])
+
+        return out
+
+    def grid_inputs(self, device, *, requires_grad=False):
+        def make_z(*args, dtype=torch.float):
+            real = torch.tensor([0, *torch.logspace(*args, dtype=torch.double)])
+            imag = torch.logspace(*args, dtype=torch.double)
+
+            # Don't test on subnormal numbers, since the AMOS code doesn't consider them
+            real[real < torch.finfo(dtype).tiny] = 0.0
+            imag[imag < torch.finfo(dtype).tiny] = 0.0
+
+            kwargs = dict(dtype=dtype, device=device, requires_grad=requires_grad)
+            real = torch.tensor(real, **kwargs)
+            imag = torch.tensor([*(-imag.flip(0)), 0, *imag], **kwargs)
+            real, imag = real[:, None], imag[None, :]
+            return torch.complex(real, imag)
+
+        out = [
+            [make_z(-350, 350, 75, dtype=torch.double)],
+            [make_z(-50, 50, 75)],
+            [make_z(-5, 5, 75, dtype=torch.double)],  # scipy is buggy on this input
+            [make_z(-5, 5, 75)],  # scipy is buggy on this input
+        ]
+        if device == "cuda":
+            # half precision is only supported on CUDA
+            out.append([make_z(-5, 5, 75, dtype=torch.half)])
+        return out
+
+    def _test_correctness(self, device):
+        samples = self.sample_inputs(device) + self.grid_inputs(device)
+        for args in samples:
+            result = torch_bessel.ops.modified_bessel_k1(*args)
+            expected = reference_modified_bessel_k1(*args)
+            if expected.dtype in {torch.float, torch.cfloat, torch.chalf}:
+                # ierr = 4, complete loss of significance
+                expected[args[0].abs() > 4194303.98419452] = torch.nan
+                # ierr = 2, overflow
+                mask = (args[0].abs() < 1.1754944e-35) & (args[0] != 0)
+                if expected.is_complex():
+                    expected[mask] = torch.nan
+                    expected[mask & (args[0].imag == 0) & (args[0] != 0)] = torch.inf
+                else:
+                    expected[mask & (args[0] != 0)] = torch.inf
+            # fix scipy bug. See https://github.com/scipy/xsf/issues/46
+            if args[0].is_complex():
+                mask = (
+                    (args[0].real > 685)
+                    & (args[0].real < 690)
+                    & (args[0].imag.abs() > 685)
+                    & (args[0].imag.abs() < 1.1e5)
+                )
+                expected[mask] = 0.0  # scipy returns incorrect/NaN output
+            torch.testing.assert_close(result, expected, equal_nan=True)
+
+    def test_correctness_cpu(self):
+        self._test_correctness("cpu")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
+    def test_correctness_cuda(self):
+        self._test_correctness("cuda")
+
+    def _opcheck(self, device):
+        # Use opcheck to check for incorrect usage of operator registration APIs
+        samples = self.sample_inputs(device, requires_grad=False)
+        for args in samples:
+            if not args[0].is_complex():
+                continue
+            opcheck(
+                torch.ops.torch_bessel.modified_bessel_k1_complex_forward.default, args
+            )
+
+    def test_opcheck_cpu(self):
+        self._opcheck("cpu")
+
+    @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires cuda")
+    def test_opcheck_cuda(self):
+        self._opcheck("cuda")
+
+
+@pytest.mark.parametrize("in_dims", [0, 1])
+def test_vmap_k1(in_dims):
+    z = torch.randn(3, 4).abs() + torch.randn(3, 4) * 1j
+    func = torch.func.vmap(torch_bessel.ops.modified_bessel_k1, in_dims=in_dims)
+    out = func(z)
+    expected = torch_bessel.ops.modified_bessel_k1(z)
+    if in_dims == 1:
+        expected = expected.t()
+    torch.testing.assert_close(out, expected)

torch_bessel-0.0.9/torch_bessel/__init__.py

@@ -0,0 +1,11 @@
+import warnings
+
+# from . import ops
+from . import _C, ops
+
+# Silence warning emitted at torch/nested/_internal/nested_tensor.py:417
+warnings.filterwarnings(
+    "ignore", message="Failed to initialize NumPy: No module named 'numpy'"
+)
+
+__all__ = ["ops"]

torch_bessel-0.0.9/torch_bessel/csrc/iterator.cpp

@@ -0,0 +1,60 @@
+// #include <torch/extension.h>
+#include <torch/library.h>
+#include <ATen/TensorIterator.h>
+
+
+at::TensorIterator build_iterator_11(const at::Tensor& result, const at::Tensor& z) {
+  return (
+    at::TensorIteratorConfig()
+    .set_check_mem_overlap(true)
+    .allow_cpu_scalars(true)
+    .enforce_safe_casting_to_output(true)
+    .add_output(result)
+    .add_input(z)
+  ).build();
+}
+
+
+at::TensorIterator build_iterator_12(const at::Tensor& result, const at::Tensor& v, const at::Tensor& z) {
+  return (
+    at::TensorIteratorConfig()
+    .set_check_mem_overlap(true)
+    .allow_cpu_scalars(true)
+    .promote_inputs_to_common_dtype(true)
+    .cast_common_dtype_to_outputs(true)
+    .enforce_safe_casting_to_output(true)
+    .promote_integer_inputs_to_float(true)
+    .add_output(result)
+    .add_input(v)
+    .add_input(z)
+  ).build();
+}
+
+at::TensorIterator build_iterator_21(const at::Tensor& result1, const at::Tensor& result2, const at::Tensor& z) {
+  return (
+    at::TensorIteratorConfig()
+    .set_check_mem_overlap(true)
+    .allow_cpu_scalars(true)
+    .enforce_safe_casting_to_output(true)
+    .add_output(result1)
+    .add_output(result2)
+    .add_input(z)
+  ).build();
+}
+
+
+at::TensorIterator build_iterator_22(const at::Tensor& result1, const at::Tensor& result2, const at::Tensor& v, const at::Tensor& z) {
+  return (
+    at::TensorIteratorConfig()
+    .set_check_mem_overlap(true)
+    .allow_cpu_scalars(true)
+    .promote_inputs_to_common_dtype(true)
+    .cast_common_dtype_to_outputs(true)
+    .enforce_safe_casting_to_output(true)
+    .promote_integer_inputs_to_float(true)
+    .add_output(result1)
+    .add_output(result2)
+    .add_input(v)
+    .add_input(z)
+  ).build();
+}

torch_bessel-0.0.9/torch_bessel/csrc/torch_bessel.cpp

@@ -0,0 +1,94 @@
+#include <Python.h>
+// #include <torch/extension.h>
+#include <torch/library.h>
+#include <ATen/ATen.h>
+#include <ATen/TensorIterator.h>
+#include <ATen/native/cpu/Loops.h>
+#include <c10/util/complex.h>
+#include <iostream>
+
+#include "bessel_k.h"
+#include "iterator.h"
+
+// See https://docs.pytorch.org/tutorials/advanced/cpp_custom_ops.html
+extern "C" {
+  /* Creates a dummy empty _C module that can be imported from Python.
+    The import from Python will load the .so consisting of this file
+    in this extension, so that the TORCH_LIBRARY static initializers
+    below are run. */
+  PyObject* PyInit__C(void)
+  {
+      static struct PyModuleDef module_def = {
+          PyModuleDef_HEAD_INIT,
+          "_C",   /* name of module */
+          NULL,   /* module documentation, may be NULL */
+          -1,     /* size of per-interpreter state of the module,
+                    or -1 if the module keeps state in global variables. */
+          NULL,   /* methods */
+      };
+      return PyModule_Create(&module_def);
+  }
+}
+
+namespace torch_bessel {
+
+at::Tensor modified_bessel_k0_complex_forward_cpu(const at::Tensor& z) {
+  TORCH_INTERNAL_ASSERT(z.device().type() == at::DeviceType::CPU);
+  at::ScalarType dtype = z.scalar_type();
+  at::Tensor result = at::empty(at::IntArrayRef(), dtype).resize_(0);
+  at::TensorIterator iter = build_iterator_11(result, z);
+  AT_DISPATCH_COMPLEX_TYPES(dtype, "modified_bessel_k0_complex_forward_cpu", [&]() {
+    at::native::cpu_kernel(iter, [](scalar_t z) -> scalar_t {
+      return modified_bessel_k0_complex_forward(z);
+    });
+  });
+  return result;
+}
+
+std::tuple<at::Tensor, at::Tensor> modified_bessel_k0_complex_forward_backward_cpu(const at::Tensor& z) {
+  TORCH_INTERNAL_ASSERT(z.device().type() == at::DeviceType::CPU);
+  at::ScalarType dtype = z.scalar_type();
+  at::Tensor result1 = at::empty(at::IntArrayRef(), dtype).resize_(0);
+  at::Tensor result2 = at::empty(at::IntArrayRef(), dtype).resize_(0);
+  at::TensorIterator iter = build_iterator_21(result1, result2, z);
+  AT_DISPATCH_COMPLEX_TYPES(dtype, "modified_bessel_k0_complex_forward_backward_cpu", [&]() {
+    at::native::cpu_kernel_multiple_outputs(iter, [](scalar_t z) -> std::tuple<scalar_t, scalar_t> {
+      scalar_t cy[2];
+      modified_bessel_k0_complex_forward_backward(z, cy);
+      return std::make_tuple(cy[0], cy[1]);
+    });
+  });
+  return std::make_tuple(result1, result2);
+}
+
+at::Tensor modified_bessel_k1_complex_forward_cpu(const at::Tensor& z) {
+  TORCH_INTERNAL_ASSERT(z.device().type() == at::DeviceType::CPU);
+  at::ScalarType dtype = z.scalar_type();
+  at::Tensor result = at::empty(at::IntArrayRef(), dtype).resize_(0);
+  at::TensorIterator iter = build_iterator_11(result, z);
+  AT_DISPATCH_COMPLEX_TYPES(dtype, "modified_bessel_k1_complex_forward_cpu", [&]() {
+    at::native::cpu_kernel(iter, [](scalar_t z) -> scalar_t {
+      return modified_bessel_k1_complex_forward(z);
+    });
+  });
+  return result;
+}
+
+// // Registers _C as a Python extension module.
+// PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}
+
+// Defines the operators
+TORCH_LIBRARY(torch_bessel, m) {
+  m.def("modified_bessel_k0_complex_forward(Tensor z) -> Tensor");
+  m.def("modified_bessel_k0_complex_forward_backward(Tensor z) -> (Tensor, Tensor)");
+  m.def("modified_bessel_k1_complex_forward(Tensor z) -> Tensor");
+}
+
+// Registers CPU implementations for bessel_k
+TORCH_LIBRARY_IMPL(torch_bessel, CPU, m) {
+  m.impl("modified_bessel_k0_complex_forward", &modified_bessel_k0_complex_forward_cpu);
+  m.impl("modified_bessel_k0_complex_forward_backward", &modified_bessel_k0_complex_forward_backward_cpu);
+  m.impl("modified_bessel_k1_complex_forward", &modified_bessel_k1_complex_forward_cpu);
+}
+
+}

torch_bessel-0.0.9/torch_bessel/ops.py

@@ -0,0 +1,134 @@
+# from pathlib import Path
+from numbers import Number
+from typing import Union
+
+import torch
+from torch import Tensor
+
+__all__ = ["modified_bessel_k0", "modified_bessel_k1"]
+
+# # load C extension before calling torch.library API, see
+# # https://pytorch.org/tutorials/advanced/cpp_custom_ops.html
+# so_dir = Path(__file__).parent
+# so_files = list(so_dir.glob("_C*.so"))
+# assert len(so_files) == 1, (
+#     f"Expected one _C*.so file at {so_dir}, found {len(so_files)}"
+# )
+# torch.ops.load_library(so_files[0])
+
+
+class ModifiedBesselK0(torch.autograd.Function):
+    @staticmethod
+    def forward(z, singularity):
+        if not z.is_complex():
+            out = (torch.special.modified_bessel_k0(z), None)
+        elif not z.requires_grad:
+            out = (
+                torch.ops.torch_bessel.modified_bessel_k0_complex_forward.default(z),
+                None,
+            )
+        else:
+            out = torch.ops.torch_bessel.modified_bessel_k0_complex_forward_backward.default(
+                z
+            )
+
+        if singularity is None:
+            return (*out, None)
+
+        mask = z != 0
+        return (out[0].where(mask, singularity), out[1], mask)
+
+    @staticmethod
+    def setup_context(ctx, inputs, outputs):
+        if ctx.needs_input_grad[1]:
+            raise NotImplementedError("Gradient w.r.t. singularity is not implemented")
+
+        if ctx.needs_input_grad[0]:
+            if outputs[1] is None:
+                ctx.save_for_backward(inputs[0], None, outputs[2])
+            else:
+                ctx.save_for_backward(None, outputs[1], outputs[2])
+
+        ctx.set_materialize_grads(False)
+
+    @staticmethod
+    def backward(ctx, grad, _, __):
+        if grad is None or not ctx.needs_input_grad[0]:
+            return (None, None)
+
+        x, deriv, mask = ctx.saved_tensors
+        if deriv is None:
+            out = -torch.special.modified_bessel_k1(x).mul_(grad)
+        else:
+            out = grad * deriv
+        if mask is not None:
+            out = out.where(mask, 0)
+        return (out, None)
+
+    @staticmethod
+    def vmap(info, in_dims, z, singularity):
+        if singularity is None and in_dims[1] is not None:
+            raise ValueError("in_dims[1] must be None if singularity is not provided.")
+
+        if in_dims[0] is not None:
+            z = z.movedim(in_dims[0], 0)
+
+        if in_dims[1] is not None:
+            singularity = singularity.movedim(in_dims[1], 0)
+
+        out = ModifiedBesselK0.apply(z, singularity)
+        out_dims = [0] * 3 if any(d is not None for d in in_dims) else [None] * 3
+        if out[1] is None:
+            out_dims[1] = None
+        if out[2] is None:
+            out_dims[2] = None
+
+        return (out, tuple(out_dims))
+
+
+def modified_bessel_k0(
+    z: Tensor, singularity: Union[Number, Tensor, None] = None
+) -> Tensor:
+    return ModifiedBesselK0.apply(z, singularity)[0]
+
+
+def modified_bessel_k1(z: Tensor) -> Tensor:
+    # non-differentiable for now
+    if not z.is_complex():
+        return torch.special.modified_bessel_k1(z)
+    return torch.ops.torch_bessel.modified_bessel_k1_complex_forward.default(z)
+
+
+@torch.library.register_fake("torch_bessel::modified_bessel_k0_complex_forward")
+def _(z):
+    return torch.empty_like(z)
+
+
+@torch.library.register_fake(
+    "torch_bessel::modified_bessel_k0_complex_forward_backward"
+)
+def _(z):
+    return torch.empty_like(z), torch.empty_like(z)
+
+
+@torch.library.register_fake("torch_bessel::modified_bessel_k1_complex_forward")
+def _(z):
+    return torch.empty_like(z)
+
+
+def modified_bessel_k0_backward(ctx, grad, _):
+    if ctx.needs_input_grad[0]:
+        return grad * ctx.saved_tensors[0]
+    return None
+
+
+def modified_bessel_k0_setup_context(ctx, inputs, output):
+    if ctx.needs_input_grad[0]:
+        ctx.save_for_backward(output[-1])
+
+
+torch.library.register_autograd(
+    "torch_bessel::modified_bessel_k0_complex_forward_backward",
+    modified_bessel_k0_backward,
+    setup_context=modified_bessel_k0_setup_context,
+)

torch_bessel-0.0.9/torch_bessel.egg-info/PKG-INFO

@@ -0,0 +1,50 @@
+Metadata-Version: 2.4
+Name: torch_bessel
+Version: 0.0.9
+Summary: PyTorch extension package for Bessel functions with arbitrary real order and complex inputs
+Home-page: https://github.com/hchau630/torch-bessel
+Author: Ho Yin Chau
+Requires-Python: >= 3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# About
+PyTorch extension package for modified Bessel functions of the second kind with complex inputs
+
+# Install
+Currently only supports Linux (with CUDA 12.4) or MacOS (Apple silicon, cpu only) with python >= 3.9, <= 3.12.
+```
+pip install torch-bessel
+```
+
+UPDATE (May 26, 2026): It seems the package might no longer be compatible with the latest pytorch version (>= 2.12.0), so try installing an earlier version of pytorch if you encounter issues.
+
+UPDATE (May 28, 2026): On Apple silicon, importing this package with pytorch version >= 2.9.0 causes a crash (bus error: 10) when exiting a python program, though the program runs normally prior to that. So consider installing torch < 2.9.0 on Apple silicon.
+
+# Example
+```
+import torch_bessel
+
+real, imag = torch.randn(2, 5, device="cuda")
+z = torch.complex(real.abs(), imag)  # correctness for inputs in the left-half complex plane is not gauranteed.
+torch_bessel.ops.modified_bessel_k0(z)
+```
+
+# Implemented functions
+- `modified_bessel_k0`: Same as `torch.special.modified_bessel_k0`, but also handles backpropagation and complex inputs on cpu and cuda. Correctness is guaranteed on the right-half complex plane. On cuda, `torch.chalf` inputs are also supported, though the underlying cuda kernel just upcasts `chalf` to `cfloat` (note that this uses no extra GPU memory, as opposed to manually casting torch.chalf to torch.cfloat before calling `modified_bessel_k0` which doubles the GPU memory used). On the left-half complex plane, function output appears mostly correct, but with small numerical errors for certain inputs. On the negative real line, output is NaN.
+- `modified_bessel_k1`: Same as `torch.special.modified_bessel_k1`, but also handles complex inputs on cpu and cuda. Backpropagation not implemented, but this can be easily manually implemented yourself by writing a torch.autograd.Function using the recurrence properties of bessel functions. Same caveats as `modified_bessel_k0` apply.
+
+# WIP
+- `modified_bessel_kv`: Analogue of `scipy.special.kv`.
+
+# Benchmarks
+Benchmarking performed with the `asv` package. Results can be viewed at https://hchau630.github.io/torch-bessel.

torch_bessel-0.0.9/torch_bessel.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+LICENSE
+README.md
+pyproject.toml
+setup.py
+benchmarks/__init__.py
+benchmarks/benchmarks.py
+tests/test_extension.py
+torch_bessel/__init__.py
+torch_bessel/ops.py
+torch_bessel.egg-info/PKG-INFO
+torch_bessel.egg-info/SOURCES.txt
+torch_bessel.egg-info/dependency_links.txt
+torch_bessel.egg-info/requires.txt
+torch_bessel.egg-info/top_level.txt
+torch_bessel/csrc/iterator.cpp
+torch_bessel/csrc/torch_bessel.cpp

torch_bessel-0.0.9/torch_bessel.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

torch_bessel-0.0.9/torch_bessel.egg-info/requires.txt

@@ -0,0 +1 @@
+torch

torch_bessel-0.0.9/torch_bessel.egg-info/top_level.txt

@@ -0,0 +1,2 @@
+benchmarks
+torch_bessel