psf_generator.utils.bessel#

A collection of custom Bessel functions with gradient tracking.

These functions contain adjoint-enabled overrides for the PyTorch build-in bessel_j0 and bessel_j1 as those do not have gradient tracking as of v1.13.1.

Classes#

`BesselJ0`	Differentiable version of PyTorch's bessel_j0(x).
`BesselJ1`	Differentiable version of bessel_j1(x).

Module Contents#

class psf_generator.utils.bessel.BesselJ0[source]#

Bases: torch.autograd.Function

Differentiable version of PyTorch’s bessel_j0(x).

static forward(ctx: Any, x: torch.Tensor) → torch.Tensor[source]#

static vjp(ctx: Any, grad_output: torch.Tensor) → torch.Tensor[source]#: Vector-Jacobian product, for reverse-mode adjoint (backward()).

static jvp(ctx: Any, grad_input: torch.Tensor) → torch.Tensor[source]#: Jacobian-vector product, for forward-mode adjoint.

class psf_generator.utils.bessel.BesselJ1[source]#

Bases: torch.autograd.Function

Differentiable version of bessel_j1(x).

static forward(ctx: Any, x: torch.Tensor) → torch.Tensor[source]#

static vjp(ctx: Any, grad_output: torch.Tensor) → torch.Tensor[source]#: Vector-Jacobian product, for reverse-mode adjoint (backward()).

static jvp(ctx: Any, grad_input: torch.Tensor) → torch.Tensor[source]#: Jacobian-vector product, for forward-mode adjoint.