Skip to content
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

Allow ChainRules zero types internally #1389

Draft
wants to merge 11 commits into
base: master
Choose a base branch
from
Draft

Allow ChainRules zero types internally #1389

wants to merge 11 commits into from

Conversation

ToucheSir
Copy link
Member

Following on #1385, this allows us to propagate ZeroTangent, NoTangent and NotImplemented inside the boundary of pullback. Doing so should prevent lossy conversions back and forth, as well as inadvertent dropping of the erroring functionality NotImplemented brings. Conversions remain the same for the user-level API, which means this only partially addresses #1227.

This PR can be seen as an incremental step towards #603. There are a couple of options for next steps:

  • Allowing Tangents to pass through without conversion
  • Supporting thunks (Utilize ChainRulesCore thunks #966)
  • Switching all internal rules to use CR zeros instead of Nothing (or converting them to rrules.

PR Checklist

  • Tests are added
  • [N/A] Documentation, if applicable

@ToucheSir ToucheSir marked this pull request as draft March 13, 2023 04:44
@ToucheSir
Copy link
Member Author

Turns out just changing Zygote's own @adjoints to support CR zeros is not enough because downstream package rules are also affected. FluxML/ZygoteRules.jl#25 should help with that and is now a blocker for this PR.

@ToucheSir ToucheSir closed this Mar 21, 2023
@ToucheSir ToucheSir reopened this Mar 21, 2023
src/compiler/chainrules.jl Outdated Show resolved Hide resolved
ToucheSir and others added 2 commits May 3, 2023 06:59
Co-authored-by: Frames White <oxinabox@ucc.asn.au>
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