torch-bessel 0.0.4__cp39-cp39-macosx_11_0_arm64.whl → 0.0.5__cp39-cp39-macosx_11_0_arm64.whl

benchmarks/benchmarks.py

@@ -36,7 +36,6 @@ class ModifiedBesselK0ForwardCPU:
         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],
@@ -50,6 +49,7 @@ class ModifiedBesselK0ForwardCUDA:
     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
     ):
@@ -77,7 +77,6 @@ class ModifiedBesselK0BackwardCPU:
         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
@@ -93,6 +92,7 @@ class ModifiedBesselK0BackwardCUDA:
         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_bessel/ops.py

@@ -65,6 +65,26 @@ class ModifiedBesselK0(torch.autograd.Function):
             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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: torch_bessel
-Version: 0.0.4
+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,12 @@ 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`.

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-39-darwin.so,sha256=38QT3Q1sn9xVivIA2qYlv9pvGLhq8vLoKIiz_AhenOk,231368
+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=lcS7ertGdTcwg6KPpE3fBtGreaKMXtC8sl4-ZeZxMCs,107
+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.4.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: cp39-cp39-macosx_11_0_arm64
 

torch_bessel-0.0.4.dist-info/RECORD

@@ -1,10 +0,0 @@
-benchmarks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-benchmarks/benchmarks.py,sha256=qRwHDxsrqRQjbExAwtWcrUtggUCL-UGJQQHnqjIoUIM,3311
-torch_bessel/_C.cpython-39-darwin.so,sha256=Nv-frnFJ35e1qBavsYzKi7hZciGLas_AloaOtrBsXiU,231368
-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=md3JO_ifs5j508p3TDNMgtQVtnQblpGEt_Wo4W56l8Y,107
-torch_bessel-0.0.4.dist-info/top_level.txt,sha256=xmyVjWSQ91kX-v8KCzl6wDwfAmbdZNWsP2EH9b9BccQ,24
-torch_bessel-0.0.4.dist-info/RECORD,,