torch-bessel 0.0.3__cp310-cp310-macosx_11_0_arm64.whl → 0.0.4__cp310-cp310-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)
+
+
+@skip_benchmark_if(not torch.cuda.is_available())
+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")
+
+    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()
+
+
+@skip_benchmark_if(not torch.cuda.is_available())
+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()
+
+    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,57 @@ 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
-
-        (x,) = ctx.saved_tensors
-        return -torch.special.modified_bessel_k1(x).mul_(grad)
-
-
-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]
+            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)
+
+
+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.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: torch_bessel
-Version: 0.0.3
+Version: 0.0.4
 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
@@ -39,3 +39,6 @@ torch_bessel.ops.modified_bessel_k0(z)
 
 # 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.4.dist-info/RECORD

@@ -0,0 +1,10 @@
+benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+benchmarks/benchmarks.py,sha256=qRwHDxsrqRQjbExAwtWcrUtggUCL-UGJQQHnqjIoUIM,3311
+torch_bessel/_C.cpython-310-darwin.so,sha256=ACQdYDXJFXSFdRL8QVsC293diLaUpKSvzLTR6K5bO4c,231080
+torch_bessel/__init__.py,sha256=oohbWz8vxekl7kqDNSWqDH3ORabf9-Tug1KJryKw51A,230
+torch_bessel/ops.py,sha256=htS41Mnz2eTcBWNCs33PSFee__S4KRfJ5zVopM2VNd8,2908
+torch_bessel-0.0.4.dist-info/LICENSE,sha256=do0DI6wu4mF3VXnEXXPYZqVEatoRSSamgz9t80wU7_o,1068
+torch_bessel-0.0.4.dist-info/METADATA,sha256=t73hRVqeL4sosAMOVmAgnJjUfOCsahFtE_jL1yHBy6Q,1347
+torch_bessel-0.0.4.dist-info/WHEEL,sha256=ezfKMaDztqf77C8lvQ0NCnZxkTaOaKLprqJ8q932MhU,109
+torch_bessel-0.0.4.dist-info/top_level.txt,sha256=xmyVjWSQ91kX-v8KCzl6wDwfAmbdZNWsP2EH9b9BccQ,24
+torch_bessel-0.0.4.dist-info/RECORD,,

{torch_bessel-0.0.3.dist-info → torch_bessel-0.0.4.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-310-darwin.so,sha256=4CEZxUFE2_LvYHrqV1bTeXn0P9LieU1eSAnRBnP6l10,231080
-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=ezfKMaDztqf77C8lvQ0NCnZxkTaOaKLprqJ8q932MhU,109
-torch_bessel-0.0.3.dist-info/top_level.txt,sha256=cbDIjTj71LuAlVyyYyDt8fOAeLaVeX3Vums5F2FBa-4,13
-torch_bessel-0.0.3.dist-info/RECORD,,