torchref.base.targets.triton.xray_ml module

Triton kernels for the Maximum-Likelihood X-ray target.

Forward computes log(I0e(arg)) using Abramowitz & Stegun 9.8.1/9.8.3 polynomial approximations split at x=3.75:

x < 3.75: I0(x) = Poly_small(t²), t = x/3.75

log(I0e(x)) = log(I0(x)) - x

x ≥ 3.75: sqrt(x)·exp(-x)·I0(x) = Poly_large(t), t = 3.75/x

log(I0e(x)) = log(Poly_large(t)) - 0.5·log(x)

Backward needs I1(x)/I0(x), which is computed from the same polynomial pair (Abramowitz 9.8.3-9.8.4). This avoids overflow at the clamped arg_bessel = 1e6 ceiling.

torchref.base.targets.triton.xray_ml.ml_xray_loss_math_triton(F_obs, F_calc, sigma, centric_flags, mask)

Triton-backed maximum-likelihood X-ray loss.

Drop-in replacement for torchref.base.targets.xray_ml.ml_xray_loss_math().