Skip to content

[AutoDiff] Initial support for differentiation of throwing functions #82653

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: main
Choose a base branch
from

Conversation

asl
Copy link
Contributor

@asl asl commented Jun 30, 2025

This adds initial support for differentiation of functions that may produce Error result. Essentially we wrap the pullback into Optional and emit a diamond-shape control flow pattern depending on whether the pullback value is available or not. VJP emission was modified to accommodate for this. In addition to this, some additional tricks are requires as try_apply result is not available in the instruction parent block, it is available in normal successor basic block.

As a result we can now:

  • differentiate an active try_apply result (that would be produced from do ... try .. catch constructions)
  • try_apply when error result is unreachable (usually try! and similar source code constructs)
  • Support (some) throwing functions with builtin differentiation operators. stdlib change will follow. Though we cannot support typed throws here (yet)
  • Correctly propagate error types during currying around differentiable functions as well as type-checking for @derivative(of:) attribute, so we can register custom derivatives for functions producing error result
  • Added custom derivative for Optional.?? operator (note that support here is not yet complete as we cannot differentiate through autoclosures, so x ?? y works only if y is not active, e.g. a constant value).

Some fixes here and there

@asl asl requested a review from rxwei June 30, 2025 22:53
@asl asl requested review from a team, hborla, slavapestov, xedin, eeckstein and jckarter as code owners June 30, 2025 22:53
@asl
Copy link
Contributor Author

asl commented Jun 30, 2025

Tagging @JaapWijnen @kovdan01

@asl
Copy link
Contributor Author

asl commented Jun 30, 2025

@swift-ci please test

@asl
Copy link
Contributor Author

asl commented Jul 1, 2025

@swift-ci please test

@asl asl closed this Jul 1, 2025
@asl asl reopened this Jul 1, 2025
@asl
Copy link
Contributor Author

asl commented Jul 1, 2025

Windows fails in unrelated IRGen/async/run-call-void-throws-to-int-throwing_call-sync-nothrow_call-async-throw.sil

@swift-ci please test windows platform

assert(inserted && "should have unique adjoint for try_apply");
break;
}
case SILValueCategory::Address: {
Copy link
Contributor

Choose a reason for hiding this comment

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

With existing tests, we never fall into this case. It's probably worth covering it with tests.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes, it seems we'd need something like object value with non-loadable tangent.

@asl
Copy link
Contributor Author

asl commented Jul 11, 2025

@swift-ci please test

@asl asl removed the request for review from a team July 11, 2025 05:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
Projects
None yet
2 participants