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

Add rules for division by Cholesky #631

Merged
merged 5 commits into from
Aug 3, 2022
Merged

Add rules for division by Cholesky #631

merged 5 commits into from
Aug 3, 2022

Conversation

sethaxen
Copy link
Member

\(::Cholesky, ::AbstractVecOrMat) and /(::AbstractMatrix, ::Cholesky) internally call ldiv! and rdiv!, respectively. This PR adds rrules for these functions to cover this up. This also allows the correspond adjoint for \ to be removed from Zygote.

Fixes #611

@sethaxen sethaxen requested a review from devmotion June 17, 2022 21:02
@devmotion
Copy link
Member

I'm not sure if this should be added here. Are there other examples in ChainRules where a rule is implemented only because the primal is mutating? At first glance this seems like a Zygote specific issue that should be addressed there. Or more generally a problem of mutation-incompatible ADs, so maybe one should only define it for RuleConfig{>:Union{ReverseMode,NoMutation}} (or something similar).

It also seems this problem might affect factorizations more generally? https://github.com/JuliaLang/julia/blob/ae89f8b00a6e785832e4f47de73599e032aa9674/stdlib/LinearAlgebra/src/factorization.jl#L103

@sethaxen
Copy link
Member Author

Are there other examples in ChainRules where a rule is implemented only because the primal is mutating?

Yes, there are number of rules defined due to mutation; most are not documented as such, but see https://github.com/JuliaDiff/ChainRules.jl/search?q=mutates. The rules don't belong in Zygote, since other ADs like Diffractor will also need them. In the past, we have added rules that cover up mutation so that they can be removed from Zygote.

Or more generally a problem of mutation-incompatible ADs, so maybe one should only define it for RuleConfig{>:Union{ReverseMode,NoMutation}} (or something similar).

Yeah that's an interesting idea. I wonder why we haven't done that in the past. @oxinabox @mcabbott do you know?

It also seems this problem might affect factorizations more generally? https://github.com/JuliaLang/julia/blob/ae89f8b00a6e785832e4f47de73599e032aa9674/stdlib/LinearAlgebra/src/factorization.jl#L103

Yes, but unfortunately the solution is specific to each factorization. We handle this one here since there's a corresponding adjoint in Zygote we'd like to remove.

@ToucheSir
Copy link
Contributor

Or more generally a problem of mutation-incompatible ADs, so maybe one should only define it for RuleConfig{>:Union{ReverseMode,NoMutation}} (or something similar).

Yeah that's an interesting idea. I wonder why we haven't done that in the past. @oxinabox @mcabbott do you know?

I suppose one hang-up is that there is no HasMutation/NoMutation type in CRC :P. Were there one, we'd update ZygoteRuleConfig to use it.

@oxinabox
Copy link
Member

 I suppose one hang-up is that there is no HasMutation/NoMutation type in CRC :P. Were there one, we'd update ZygoteRuleConfig to use it.

I am in no hurry to add it.
Later we can, but first i want to see a real AD that has CR support and handles mutatation well.
Yota also doesn't.
And nor does Nabla.

and ReverseDiff is opt in per rule anyway,

@github-actions github-actions bot added the needs version bump Version needs to be incremented or set to -DEV in Project.toml label Aug 3, 2022
@github-actions github-actions bot removed the needs version bump Version needs to be incremented or set to -DEV in Project.toml label Aug 3, 2022
@sethaxen sethaxen merged commit 8587a07 into main Aug 3, 2022
@sethaxen sethaxen deleted the choleskyldiv branch August 3, 2022 11:00
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.

Wrong call on cholesky rrule
4 participants