torch-bessel 0.0.6__cp312-cp312-macosx_11_0_arm64.whl → 0.0.7__cp312-cp312-macosx_11_0_arm64.whl

benchmarks/benchmarks.py

@@ -17,6 +17,17 @@ def _setup(
     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],
@@ -97,3 +108,36 @@ class ModifiedBesselK0BackwardCUDA:
         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/ops.py

@@ -5,7 +5,7 @@ from typing import Union
 import torch
 from torch import Tensor
 
-__all__ = ["modified_bessel_k0"]
+__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
@@ -92,6 +92,13 @@ def modified_bessel_k0(
     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)
@@ -104,6 +111,11 @@ 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]

{torch_bessel-0.0.6.dist-info → torch_bessel-0.0.7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torch_bessel
-Version: 0.0.6
+Version: 0.0.7
 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
@@ -18,7 +18,7 @@ Dynamic: requires-python
 Dynamic: summary
 
 # About
-PyTorch extension package for Bessel functions with arbitrary real order and complex inputs
+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.
@@ -37,6 +37,7 @@ 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 for double types, and almost guaranteed for float types, though it appears there are a very small handful of inputs which result in NaNs which needs to be fixed. 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`.

torch_bessel-0.0.7.dist-info/RECORD

@@ -0,0 +1,10 @@
+benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+benchmarks/benchmarks.py,sha256=_Bi3QFknID_n99lrqMTW36YFhCEIxxQzu5hmDd1inkY,4595
+torch_bessel/_C.cpython-312-darwin.so,sha256=0SvorSYRDqUrLH5GZMgWQ50F0-o6IggJQxjjcAtPQy8,233368
+torch_bessel/__init__.py,sha256=oohbWz8vxekl7kqDNSWqDH3ORabf9-Tug1KJryKw51A,230
+torch_bessel/ops.py,sha256=U-o-3brRTM4UX8cJqZInjsf2zLSaGzGX2yUZPYpq19s,3968
+torch_bessel-0.0.7.dist-info/licenses/LICENSE,sha256=do0DI6wu4mF3VXnEXXPYZqVEatoRSSamgz9t80wU7_o,1068
+torch_bessel-0.0.7.dist-info/METADATA,sha256=KG8_9O1eMuktubsS24MLnFBAk9xGKmsErvH24BuxAYQ,2374
+torch_bessel-0.0.7.dist-info/WHEEL,sha256=CltXN3lQvXbHxKDtiDwW0RNzF8s2WyBuPbOAX_ZeQlA,109
+torch_bessel-0.0.7.dist-info/top_level.txt,sha256=xmyVjWSQ91kX-v8KCzl6wDwfAmbdZNWsP2EH9b9BccQ,24
+torch_bessel-0.0.7.dist-info/RECORD,,

torch_bessel-0.0.6.dist-info/RECORD

@@ -1,10 +0,0 @@
-benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-benchmarks/benchmarks.py,sha256=LZDGAlKEWwo-VCenNTGn98bJzj9FSlQbFK5pPXRvDyI,3319
-torch_bessel/_C.cpython-312-darwin.so,sha256=YZJPORMvIiAa-JColLS3KoyD0HI60ow1aRyuzjlsif8,232280
-torch_bessel/__init__.py,sha256=oohbWz8vxekl7kqDNSWqDH3ORabf9-Tug1KJryKw51A,230
-torch_bessel/ops.py,sha256=P4qo1GN6XExuKfI061dsZMpAPuimn6gYt-PwRmR4_Fk,3584
-torch_bessel-0.0.6.dist-info/licenses/LICENSE,sha256=do0DI6wu4mF3VXnEXXPYZqVEatoRSSamgz9t80wU7_o,1068
-torch_bessel-0.0.6.dist-info/METADATA,sha256=uViUBBZq-V6HuSKg2N4NMiWJYtJwUxmLxkB0dtTVzS4,2034
-torch_bessel-0.0.6.dist-info/WHEEL,sha256=CltXN3lQvXbHxKDtiDwW0RNzF8s2WyBuPbOAX_ZeQlA,109
-torch_bessel-0.0.6.dist-info/top_level.txt,sha256=xmyVjWSQ91kX-v8KCzl6wDwfAmbdZNWsP2EH9b9BccQ,24
-torch_bessel-0.0.6.dist-info/RECORD,,