Skip to content

Remove TorchScript helpers from linear CG - #134

Merged
kayween merged 1 commit into
mainfrom
remove-linear-cg-jit
Sep 11, 2026
Merged

kayween merged 1 commit into
mainfrom
remove-linear-cg-jit

Conversation

@saitcakmak

Copy link
Copy Markdown
Collaborator

Summary

  • remove the deprecated torch.jit.script helpers from linear_cg
  • consolidate preconditioned and unpreconditioned iterations into one code path
  • add direct coverage for a custom Jacobi preconditioner

Why

The helpers were introduced as TorchScript-only optimization boundaries. Replacing them with torch.compile, as proposed in #129, adds substantial first-call compilation overhead and shape-specialization costs for small fragments inside a Python loop. The default preconditioner already implements the unpreconditioned behavior by cloning the residual, so both branches can share the same update code without either compilation mechanism.

This also removes the torch.jit.script deprecation warning emitted when importing linear_operator on recent PyTorch versions.

Correctness

The consolidated loop preserves the existing CG update order. A new float64 parity test compares a Jacobi-preconditioned solve with torch.linalg.solve and verifies that the supplied preconditioner is used repeatedly. Temporarily reversing the residual-update sign caused this test to fail, confirming that it detects an incorrect update.

Test plan

  • ufmt on both changed files
  • python -m flake8 -j 1 linear_operator/utils/linear_cg.py test/utils/test_linear_cg.py
  • python -m pytest test/utils/test_linear_cg.py test/functions/test_inv_quad_logdet.py test/operators/test_added_diag_linear_operator.py test/operators/test_low_rank_root_added_diag_linear_operator.py -xvs (403 passed)
  • python -Wd -c "import linear_operator" (no TorchScript deprecation warning)

@kayween kayween left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It seems good to me to remove torch.jit.script given that it's already deprecated a long time ago.

@kayween
kayween merged commit 1c5e24b into main Sep 11, 2026
4 checks passed
@saitcakmak
saitcakmak deleted the remove-linear-cg-jit branch September 11, 2026 14:45
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants