torch-bessel 0.0.3__cp312-cp312-macosx_11_0_arm64.whl → 0.0.5__cp312-cp312-macosx_11_0_arm64.whl

benchmarks/benchmarks.py

@@ -0,0 +1,99 @@
+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
+
+
+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()

torch_bessel/ops.py

@@ -1,4 +1,6 @@
 from pathlib import Path
+from numbers import Number
+from typing import Union
 
 import torch
 from torch import Tensor
@@ -17,32 +19,77 @@ torch.ops.load_library(so_files[0])
 
 class ModifiedBesselK0(torch.autograd.Function):
     @staticmethod
-    def forward(x):
-        return torch.special.modified_bessel_k0(x)
+    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, _):
+    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]:
-            ctx.save_for_backward(*inputs)
+            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):
+    def backward(ctx, grad, _, __):
         if grad is None or not ctx.needs_input_grad[0]:
-            return None
+            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
 
-        (x,) = ctx.saved_tensors
-        return -torch.special.modified_bessel_k1(x).mul_(grad)
+        return (out, tuple(out_dims))
 
 
-def modified_bessel_k0(z: Tensor) -> Tensor:
-    if not z.is_complex():
-        return ModifiedBesselK0.apply(z)
-    if not z.requires_grad:
-        return torch.ops.torch_bessel.modified_bessel_k0_complex_forward.default(z)
-    return torch.ops.torch_bessel.modified_bessel_k0_complex_forward_backward.default(
-        z
-    )[0]
+def modified_bessel_k0(
+    z: Tensor, singularity: Union[Number, Tensor, None] = None
+) -> Tensor:
+    return ModifiedBesselK0.apply(z, singularity)[0]
 
 
 @torch.library.register_fake("torch_bessel::modified_bessel_k0_complex_forward")

{torch_bessel-0.0.3.dist-info → torch_bessel-0.0.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: torch_bessel
-Version: 0.0.3
+Version: 0.0.5
 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
@@ -30,12 +30,15 @@ pip install torch-bessel
 import torch_bessel
 
 real, imag = torch.randn(2, 5, device="cuda")
-z = torch.complex(real.abs(), imag)  # inputs on the left-half complex plane are set to NaNs.
+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 with $\mathrm{Re}(z) \geq 0$ on cpu and cuda.
+- `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 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.
 
 # 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.5.dist-info/RECORD

@@ -0,0 +1,10 @@
+benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+benchmarks/benchmarks.py,sha256=LZDGAlKEWwo-VCenNTGn98bJzj9FSlQbFK5pPXRvDyI,3319
+torch_bessel/_C.cpython-312-darwin.so,sha256=ARpLpmh0MCoSGplMTVRxOWtT2zCGI8DbJDY7satv5NM,232280
+torch_bessel/__init__.py,sha256=oohbWz8vxekl7kqDNSWqDH3ORabf9-Tug1KJryKw51A,230
+torch_bessel/ops.py,sha256=P4qo1GN6XExuKfI061dsZMpAPuimn6gYt-PwRmR4_Fk,3584
+torch_bessel-0.0.5.dist-info/LICENSE,sha256=do0DI6wu4mF3VXnEXXPYZqVEatoRSSamgz9t80wU7_o,1068
+torch_bessel-0.0.5.dist-info/METADATA,sha256=LX6X8j53vJCYKaXLfjlk3X_RwRYHlkoQVYC7CSJJ1vc,1561
+torch_bessel-0.0.5.dist-info/WHEEL,sha256=QEo1-fvjBiv4iPfIbG-kA4GNtCNK6HOxclrG5aWVgHI,109
+torch_bessel-0.0.5.dist-info/top_level.txt,sha256=xmyVjWSQ91kX-v8KCzl6wDwfAmbdZNWsP2EH9b9BccQ,24
+torch_bessel-0.0.5.dist-info/RECORD,,

{torch_bessel-0.0.3.dist-info → torch_bessel-0.0.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.0)
+Generator: setuptools (76.0.0)
 Root-Is-Purelib: false
 Tag: cp312-cp312-macosx_11_0_arm64
 

{torch_bessel-0.0.3.dist-info → torch_bessel-0.0.5.dist-info}/top_level.txt

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

torch_bessel-0.0.3.dist-info/RECORD

@@ -1,8 +0,0 @@
-torch_bessel/_C.cpython-312-darwin.so,sha256=d2qpLB8simaybQxmgXOrrHwvsG3yw32yoehTIX5PsO0,232280
-torch_bessel/__init__.py,sha256=oohbWz8vxekl7kqDNSWqDH3ORabf9-Tug1KJryKw51A,230
-torch_bessel/ops.py,sha256=Q9BrLxi15MS53xSt_S9dyE3g8_8_GCFYhfAztIor8Fw,2043
-torch_bessel-0.0.3.dist-info/LICENSE,sha256=do0DI6wu4mF3VXnEXXPYZqVEatoRSSamgz9t80wU7_o,1068
-torch_bessel-0.0.3.dist-info/METADATA,sha256=7hH6thd8bu4oaKYfNq5v6V4JkUe6g6_C9Irl618y3-8,1220
-torch_bessel-0.0.3.dist-info/WHEEL,sha256=VujM3ypTCyUW6hcTDdK2ej0ARVMxlU1Djlh_zWnDgqk,109
-torch_bessel-0.0.3.dist-info/top_level.txt,sha256=cbDIjTj71LuAlVyyYyDt8fOAeLaVeX3Vums5F2FBa-4,13
-torch_bessel-0.0.3.dist-info/RECORD,,