torch-bessel 0.0.1__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

torch_bessel/__init__.py

@@ -0,0 +1,9 @@
+import warnings
+from . import 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/ops.py

@@ -0,0 +1,75 @@
+from pathlib import Path
+
+import torch
+from torch import Tensor
+
+__all__ = ["modified_bessel_k0"]
+
+# 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(x):
+        return torch.special.modified_bessel_k0(x)
+
+    @staticmethod
+    def setup_context(ctx, inputs, _):
+        if ctx.needs_input_grad[0]:
+            ctx.save_for_backward(*inputs)
+        ctx.set_materialize_grads(False)
+
+    @staticmethod
+    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]
+
+
+@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)
+
+
+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.1.dist-info/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.1.dist-info/METADATA

@@ -0,0 +1,39 @@
+Metadata-Version: 2.2
+Name: torch_bessel
+Version: 0.0.1
+Summary: PyTorch extension package for Bessel functions with arbitrary real order and complex inputs
+Home-page: https://github.com/pytorch/torch-bessel
+Author: Ho Yin Chau
+Requires-Python: >= 3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# About
+PyTorch extension package for Bessel functions with arbitrary real order and complex inputs
+
+# Install
+```
+pip install torch-bessel
+```
+
+# Example usage
+```
+import torch_bessel
+
+z = torch.randn(10) + 1j
+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{Im}(z) \geq 0$ on cpu and cuda.
+
+# WIP
+- `modified_bessel_kv`: Analogue of `scipy.special.kv`.

torch_bessel-0.0.1.dist-info/RECORD

@@ -0,0 +1,8 @@
+torch_bessel-0.0.1.dist-info/METADATA,sha256=HZCm9quNFZY457PtJREAYbwLamgZcp-7DkJrX8kO0xY,999
+torch_bessel-0.0.1.dist-info/RECORD,,
+torch_bessel-0.0.1.dist-info/top_level.txt,sha256=cbDIjTj71LuAlVyyYyDt8fOAeLaVeX3Vums5F2FBa-4,13
+torch_bessel-0.0.1.dist-info/WHEEL,sha256=bWgrYGwmr1oqID2askDfuFruhD_BtRXYD33vDqPTClc,151
+torch_bessel-0.0.1.dist-info/LICENSE,sha256=do0DI6wu4mF3VXnEXXPYZqVEatoRSSamgz9t80wU7_o,1068
+torch_bessel/_C.cpython-311-x86_64-linux-gnu.so,sha256=VxOjh_0ychPr4TT96LKvcOyNzmeTY3v9JBn4kfABOh8,14783376
+torch_bessel/ops.py,sha256=Q9BrLxi15MS53xSt_S9dyE3g8_8_GCFYhfAztIor8Fw,2043
+torch_bessel/__init__.py,sha256=oohbWz8vxekl7kqDNSWqDH3ORabf9-Tug1KJryKw51A,230

torch_bessel-0.0.1.dist-info/WHEEL

@@ -0,0 +1,6 @@
+Wheel-Version: 1.0
+Generator: setuptools (75.8.0)
+Root-Is-Purelib: false
+Tag: cp311-cp311-manylinux_2_17_x86_64
+Tag: cp311-cp311-manylinux2014_x86_64
+

torch_bessel-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+torch_bessel