*_solve methods: implements forward AD#65546
Closed
nikitaved wants to merge 18 commits into
Closed
Conversation
CI Flow Status⚛️ CI FlowRuleset - Version:
You can add a comment to the PR and tag @pytorchbot with the following commands: # ciflow rerun, "ciflow/default" will always be added automatically
@pytorchbot ciflow rerun
# ciflow rerun with additional labels "-l <ciflow/label_name>", which is equivalent to adding these labels manually and trigger the rerun
@pytorchbot ciflow rerun -l ciflow/scheduled -l ciflow/slowFor more information, please take a look at the CI Flow Wiki. |
Contributor
🔗 Helpful links
💊 CI failures summary and remediationsAs of commit 99f0767 (more details on the Dr. CI page): ✅ None of the CI failures appear to be your fault 💚
1 job timed out:
❄️ 1 failure tentatively classified as flakybut reruns have not yet been triggered to confirm:
|
nikitaved
commented
Sep 23, 2021
lezcano
approved these changes
Sep 24, 2021
Collaborator
lezcano
left a comment
There was a problem hiding this comment.
This is a great PR! I had thought about implementing something like that, but I never got around to do it.
The code is very clean, thanks!
| scalar_t* A_working_ptr = &A_data[i * A_mat_stride]; | ||
| scalar_t* b_working_ptr = &b_data[i * b_mat_stride]; | ||
| lapackCholeskySolve<scalar_t>(uplo, n, nrhs, A_working_ptr, n, b_working_ptr, n, &info); | ||
| lapackCholeskySolve<scalar_t>(uplo, n, nrhs, A_working_ptr, ldab, b_working_ptr, ldab, &info); |
| const bool unitriangular | ||
| ) { | ||
| return generic_solve_jvp( | ||
| [=](const Tensor& A, const Tensor& B) { |
Collaborator
There was a problem hiding this comment.
Suggested change
| [=](const Tensor& A, const Tensor& B) { | |
| [upper, transpose, unitriangular](const Tensor& A, const Tensor& B) { |
Otherwise it might need to copy all the tensors (I don't know if the compiler is allowed to ellide the copies).
Same below.
Collaborator
Author
There was a problem hiding this comment.
Not sure it will happen either way since afaik Tensors are not being copied on operator=.
Collaborator
There was a problem hiding this comment.
It was more about the copies of tensors and the reference bumps. A very very minor thing really.
Collaborator
There was a problem hiding this comment.
You can also capture everything by reference because the lambda never outlives the current scope?
14 tasks
albanD
reviewed
Sep 24, 2021
Collaborator
albanD
left a comment
There was a problem hiding this comment.
Very nice refactoring! Thanks!
| const bool unitriangular | ||
| ) { | ||
| return generic_solve_jvp( | ||
| [=](const Tensor& A, const Tensor& B) { |
Collaborator
There was a problem hiding this comment.
You can also capture everything by reference because the lambda never outlives the current scope?
…taved/solve_methods_forward_AD
albanD
approved these changes
Sep 27, 2021
Collaborator
albanD
left a comment
There was a problem hiding this comment.
Looks good.
Thanks for the updates.
Contributor
|
@albanD has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
Collaborator
|
Unlanding as this may have broken periodic_pytorch_xenial_cuda10_2_cudnn7_gcc7_old_gradcheck_test1. Relevant snippet: We had better run the periodic gradcheck job on this before relanding. |
…taved/solve_methods_forward_AD
Collaborator
|
Here is a link to the failing CI: https://app.circleci.com/pipelines/github/pytorch/pytorch/386383/workflows/8e2b4619-efd2-44dd-bd5f-24a1289b65b5/jobs/16239595 |
Collaborator
Author
|
Ah, good point, @IvanYashchuk ! That makes sense, thank you! |
Collaborator
|
That periodic job runs gradchecks with fast_mode=False. You could try testing it locally by adding |
Collaborator
|
with this label, we should have the |
…taved/solve_methods_forward_AD
…om/pytorch/pytorch into nikitaved/solve_methods_forward_AD
Collaborator
Author
|
This one looks good now with the updated OpInfo. |
albanD
approved these changes
Oct 6, 2021
Collaborator
albanD
left a comment
There was a problem hiding this comment.
Thanks for the update!
Try 2!
Contributor
|
@albanD has imported this pull request. If you are a Facebook employee, you can view this diff on Phabricator. |
laurentdupin
pushed a commit
to laurentdupin/pytorch
that referenced
this pull request
Apr 25, 2026
Summary: This PR adds forward AD for `*_solve` methods. Additionally, `cholesky_solve` gets OpInfo + a bug fix when wrong leading dimensions could be passed to LAPACK, and `lu_solve` gets forward AD with 2x`lu_solve` instead of 1x`lu_solve` + 2x`triangular_solve`. cc ezyang albanD zou3519 gqchen pearu nikitaved soulitzer Lezcano Varal7 jianyuh mruberry walterddr IvanYashchuk xwang233 Pull Request resolved: pytorch#65546 Reviewed By: gchanan Differential Revision: D31206837 Pulled By: albanD fbshipit-source-id: 040beda97442e7a88a9df9abc7bb18313ce55bc3
laurentdupin
pushed a commit
to laurentdupin/pytorch
that referenced
this pull request
Apr 25, 2026
Summary: This PR adds forward AD for `*_solve` methods. Additionally, `cholesky_solve` gets OpInfo + a bug fix when wrong leading dimensions could be passed to LAPACK, and `lu_solve` gets forward AD with 2x`lu_solve` instead of 1x`lu_solve` + 2x`triangular_solve`. cc ezyang albanD zou3519 gqchen pearu nikitaved soulitzer Lezcano Varal7 jianyuh mruberry walterddr IvanYashchuk xwang233 Pull Request resolved: pytorch#65546 Reviewed By: dagitses Differential Revision: D31431847 Pulled By: albanD fbshipit-source-id: 0e343e0d9da3c3d2051fca215fad289d77275251
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
7 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR adds forward AD for
*_solvemethods.Additionally,
cholesky_solvegets OpInfo + a bug fix when wrong leading dimensions could be passed to LAPACK,and
lu_solvegets forward AD with 2xlu_solveinstead of 1xlu_solve+ 2xtriangular_solve.cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7 @jianyuh @mruberry @walterddr @IvanYashchuk @xwang233