cholesky_solve_backward: speed up using output_mask

[mlubin]

Introduces a faster path for cholesky_solve_backward when the gradient with respect to the cholesky factor isn't required.

Adds test coverage in test_linalg.py.

Example

Setup

import torch
torch.set_num_threads(1)
mat = torch.randn(500, 1000)
mat = mat @ mat.T
L = torch.linalg.cholesky(mat, upper=False)

rhs = torch.randn(500, 1)
rhs.requires_grad = True

sol = torch.cholesky_solve(rhs, L, upper=False).sum(dim=0)

Before

%timeit torch.autograd.grad(sol, rhs, retain_graph=True)
2.61 ms ± 18.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

After

%timeit torch.autograd.grad(sol, rhs, retain_graph=True)
109 µs ± 3.42 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/112981

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ You can merge normally! (2 Unrelated Failures)

As of commit 8d3e50a with merge base db7a3cc ():

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• pull / linux-focal-py3.11-clang10 / test (dynamo, 1, 2, linux.2xlarge) (gh)
• pull / linux-focal-py3.8-clang10 / test (dynamo, 1, 2, linux.2xlarge) (gh)

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: mlubin / name: Miles Lubin (361b48e, 94a3bd8, 352a51d, 736e269, 8d3e50a)

[albanD]

The change sounds good thanks!
@lezcano where are the autograd tests for this?

[lezcano]

There is an op-info at common_methods_invocations.py, but that will not test this new route. I would suggest you add a one-off test in test_linalg.py

[lezcano]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

[pytorchmergebot]

Merge failed

Reason: 2 mandatory check(s) failed. The first few are:

• pull
• Lint

Dig deeper by viewing the failures on hud

Details for Dev Infra team

Raised by workflow job

Failing merge rule: Core Maintainers

[lezcano]

Yep, I also thought about that nit with upper=False. Thank you for fixing it!

[lezcano]

@pytorchbot merge

[pytorchmergebot]

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here