Context
We noticed pygam's PIRLS loop spends the large majority of its wall-clock time in a single
numpy.linalg.qr(WB.toarray()) call (pygam.py:783). Profiling a real workload — LogisticGAM
with a monotonic-constrained TensorTerm, fit under thread-pinned BLAS (OMP_NUM_THREADS=1 to have
one process per core in a multi-process worker pool) — showed this call dominating fit time.
What we tried
We monkeypatched this call site to route through torch.linalg.qr instead of
numpy.linalg.qr (converting the dense array to a CPU tensor and back), keeping everything
else about the fit identical. On real data (10 representative fits, various data sizes),
this consistently cut per-fit wall time by roughly 2.3x (mean 8.72s -> 3.74s across the
sample), with the effect holding across every fit we tested, not just on average.
We verified correctness by comparing predicted values between the numpy-QR and torch-QR runs:
max absolute difference across ~1000 column comparisons was ~1.3e-13.
Why this might matter for other users
This seems like it could generalize beyond our specific use case: any model with enough
observations/spline terms that WB.toarray() produces a reasonably large dense matrix would hit
the same QR-dominated cost, and any user running many concurrent fits under thread-pinned BLAS
(a common setup to avoid oversubscription in multiprocessing pools) would likely see a similar
win, since torch's linear algebra backend appears to handle this workload more efficiently in
that configuration.
Proposal
pygam already has precedent for exactly this kind of optional, auto-detected accelerated
backend -- pygam/utils.py's cholesky() function checks if SKSPIMPORT and transparently
uses sksparse.cholmod when available, falling back to scipy.linalg.cholesky otherwise. A
similar pattern could apply here: if torch is importable, use torch.linalg.qr for the PIRLS
QR step; otherwise fall back to the current numpy.linalg.qr path. This would keep torch as a
fully optional dependency (no behavior change for anyone who doesn't have it installed) while
giving users who do have it a meaningful, opt-in speedup.
Context
We noticed pygam's PIRLS loop spends the large majority of its wall-clock time in a single
numpy.linalg.qr(WB.toarray())call (pygam.py:783). Profiling a real workload — LogisticGAMwith a monotonic-constrained TensorTerm, fit under thread-pinned BLAS (
OMP_NUM_THREADS=1to haveone process per core in a multi-process worker pool) — showed this call dominating fit time.
What we tried
We monkeypatched this call site to route through
torch.linalg.qrinstead ofnumpy.linalg.qr(converting the dense array to a CPU tensor and back), keeping everythingelse about the fit identical. On real data (10 representative fits, various data sizes),
this consistently cut per-fit wall time by roughly 2.3x (mean 8.72s -> 3.74s across the
sample), with the effect holding across every fit we tested, not just on average.
We verified correctness by comparing predicted values between the numpy-QR and torch-QR runs:
max absolute difference across ~1000 column comparisons was ~1.3e-13.
Why this might matter for other users
This seems like it could generalize beyond our specific use case: any model with enough
observations/spline terms that
WB.toarray()produces a reasonably large dense matrix would hitthe same QR-dominated cost, and any user running many concurrent fits under thread-pinned BLAS
(a common setup to avoid oversubscription in multiprocessing pools) would likely see a similar
win, since torch's linear algebra backend appears to handle this workload more efficiently in
that configuration.
Proposal
pygam already has precedent for exactly this kind of optional, auto-detected accelerated
backend --
pygam/utils.py'scholesky()function checksif SKSPIMPORTand transparentlyuses
sksparse.cholmodwhen available, falling back toscipy.linalg.choleskyotherwise. Asimilar pattern could apply here: if
torchis importable, usetorch.linalg.qrfor the PIRLSQR step; otherwise fall back to the current
numpy.linalg.qrpath. This would keep torch as afully optional dependency (no behavior change for anyone who doesn't have it installed) while
giving users who do have it a meaningful, opt-in speedup.