Test changes from https://github.com/FluxML/ZygoteRules.jl/pull/26

[ToucheSir]

We have both Zygote's own CI and downstream tests here, so it'll be much easier to see what the impact is.

PR Checklist

• Tests are added
• Documentation, if applicable

[devmotion]

Seems the Zygote test failures only point to a few tests that have to be updated since they rely on the current non-collapsing behaviour.

[ToucheSir]

Most if not all of those should be fixed by #1328, but it hasn't been merged yet. Do you have any insight into the DynamicPPL error? If it doesn't look related, the ZygoteRules PR should be good to go.

[devmotion]

DynamicPPL can be ignored - AbstractMCMC 4.5 contains a breaking change which was accidentally released in a non-breaking release (TuringLang/AbstractMCMC.jl#132 (comment)). Hopefully it will be yanked soon (JuliaRegistries/General#94031).

[devmotion]

I went through the CI logs once more and I think https://github.com/FluxML/Zygote.jl/actions/runs/6631670834/job/18015683033?pr=1466#step:6:818 might be problematic: There (nothing,) seems to be correct - at least I assumed the return value of gradient should always be a Tuple with one element for each of the function arguments. Otherwise splatting, indexing, etc. could lead to nasty surprises. I still think that the other two failing tests are actually improved by the ZygoteRules change because tuples of nothings that are returned by pullbacks are simplified.

[devmotion]

I guess my assumption about the return type of gradient was wrong (but this is not documented anywhere, is it?):

Zygote.jl/src/compiler/interface.jl

Line 98 in f755127

isnothing(grad) ? nothing : map(_project, args, grad)

[ToucheSir]

Yeah, so this goes all the way back to the early early days of Zygote and I don't think the behaviour nor the reasoning for it have ever been documented: https://github.com/FluxML/Zygote.jl/blame/087b55467afb4ee312ad1325de02bdf91e07e790/src/compiler/interface.jl#L60. We should make a more formal guarantee now, but of what? It would be good to know where other AD libraries fall on this, AbstractDifferentiation.jl and its backends come to mind. Also cc @mcabbott since I think he last had to deal with this part of the codebase.

[devmotion]

Yeah, I think such conventions should be documented more clearly and followed more consistently. IMO it would be reasonable to demand that gradient always returns a tuple with gradients for each input argument (as done e.g. in AbstractDifferentiation). Also for the pullback functions, I guess it would be helpful for developers to document that nothings should always be collapsed.