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

remove @adjoint function cholesky #1114

Merged
merged 13 commits into from
Jun 20, 2022

Conversation

st--
Copy link
Contributor

@st-- st-- commented Nov 4, 2021

The adjoint of cholesky is already defined in ChainRules.jl, but it is overridden by Zygote's definition, which is missing the case for Hermitian:

C = rand(5,5)
f(x) = sum(cholesky(Hermitian(x'x)).L)
Zygote.gradient(f, C)  # errors

This PR bumps the julia compat to 1.6 and ChainRules to 1.35.3, at which point we can safely remove the adjoint. It also adds a test for the Hermitian case.

@DhairyaLGandhi
Copy link
Member

I'm generally on board. The integration test failures are related to DistributionsAD but the buildkite errors are curious seeing as they shouldn't be involved here at all.

@st--
Copy link
Contributor Author

st-- commented Nov 4, 2021

Yeah, this failure in the CI / Julia 1 test looks like something is actually not going quite right in the cholesky gradient: https://github.com/FluxML/Zygote.jl/runs/4106058861?check_suite_focus=true#step:6:189

@devmotion
Copy link
Collaborator

Maybe one should check if #1104 fixed (some of) the issues and rerun the tests based on the master branch?

@st--
Copy link
Contributor Author

st-- commented Nov 9, 2021

Maybe one should check if #1104 fixed (some of) the issues and rerun the tests based on the master branch?

There's 28 more tests in total and 3 more broken tests, comparing master-merged-in vs previous run... so doesn't seem to have helped unfortunately...

@devmotion
Copy link
Collaborator

I just remembered JuliaDiff/ChainRules.jl#611, I think it causes most of the test errors here (JuliaDiff/ChainRules.jl#611 (comment)).

@devmotion
Copy link
Collaborator

The problems with the CR rule for cholesky were fixed in CR 1.35.3, so it should be possible to rerun the tests and finish the PR by adding a lower bound for ChainRules 1.35.3.

@devmotion
Copy link
Collaborator

Looks great 👍

The remaining test error is caused by a newly added feature of the rule in ChainRules - it also supports failed factorizations for some input types now. So updating the test should be sufficient I assume.

ChainRules 1.35.3 also only supports Julia >= 1.6 (the Julia compat in CR was updated quite a while ago), so the Julia compat should be updated or the Zygote rule only be removed on Julia >= 1.6 if you want to keep support for older Julia versions.

The downstream test error of DynamicPPL was fixed in the latest release of DistributionsAD (just released around 1-2 hours ago).

@ToucheSir
Copy link
Member

@FluxML/ad @CarloLucibello should we take the opportunity to bump compat to 1.6? Both upstream (CR) and downstream (Flux) have already done so now.

@CarloLucibello
Copy link
Member

Yes let's bump

@st--
Copy link
Contributor Author

st-- commented Jun 18, 2022

@devmotion I've tried to think about how to fix the test but I couldn't think of anything, so I've just removed it for now; if you've got any better idea what a test should look like please just suggest

Copy link
Collaborator

@devmotion devmotion left a comment

Choose a reason for hiding this comment

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

Can you also update the Github action to 1.6?

test/gradcheck.jl Outdated Show resolved Hide resolved
@ToucheSir
Copy link
Member

If this looks good, we can roll it and #1226 into v0.6.41.

@ToucheSir
Copy link
Member

I'll take the +1 as approval from @devmotion, so if you wouldn't mind rebasing @st-- we can get this merged.

Copy link
Collaborator

@devmotion devmotion left a comment

Choose a reason for hiding this comment

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

A more formal approval 🙂 Looks good to me, tests seem to pass now (apart from nightlies and NeuralPDE which fail on master as well).

@ToucheSir
Copy link
Member

Molly.jl failure is either spurious or caused by changes outside this PR, will see what CI says post-merge. Always a good day when another @adjoint bites the dust :)

@ToucheSir ToucheSir merged commit c4b4fa9 into FluxML:master Jun 20, 2022
@st-- st-- deleted the st/remove_cholesky_adjoint branch June 20, 2022 20:43
@mcabbott mcabbott added the ChainRules adjoint -> rrule, and further integration label Jul 4, 2022
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ChainRules adjoint -> rrule, and further integration
Projects
None yet
Development

Successfully merging this pull request may close these issues.

6 participants