Adjoint for parent for LowerTriangular and UpperTriangular

[torfjelde]

LowerTriangular and UpperTriangular is missing adjoints for parent which can lead to downstream issues: TuringLang/Bijectors.jl#280 (comment)

This PR adds those, similar to the existing for Adjoint and Transposed.

It's unclear to me whether this should actually go into ChainRules.jl or not 😕

I'll add tests in a bit.

[CarloLucibello]

Ideally all of them should be moved to ChainRules, but that can be done in a future PR. Can you add a test?

[torfjelde]

Added tests 👍

[torfjelde]

Are the test failures related to this PR?

[ToucheSir]

They do not appear to be, but we should fix them. Unfortunately I think that requires some coordination on the AbstractFFTs side and I have no idea what's going on there or with the FFT rules here. Perhaps @gaurav-arya or someone else in the know could help?

[gaurav-arya]

Looks to be erroring on the FFT Chain Rules that have existed for quite a while in AbstractFFTs.jl, when an object of type Zygote.OneElement is provided to the pullback. The error seems to occur in the convert line added in JuliaMath/AbstractFFTs.jl#105.

Are objects of type Zygote.OneElement meant to be passed into a pullback, and how are they meant to be handled? I just tried to play a bit locally, and the object didn't seem to support many operations, e.g. broadcasting or copying or scaling.

[gaurav-arya]

OK, I see that Zygote.OneElement does indeed support those things, I just messed up by constructing its axes with unit ranges rather than Base.OneTo's.

So the error is indeed due to the convert line introduced in JuliaMath/AbstractFFTs.jl#105.

[torfjelde]

Btw what's the release schedule like for Zygote? Would it be possible to make a new release somewhat soonish?:)

[ToucheSir]

Was hoping to cut a patch release after merging this PR, but I need to figure out whether the current failures on CI (ref. #1446) require a fix on the Zygote side to ensure we're not breaking anyone's code. However I am not the best person to investigate this since I have little knowledge of the interaction between the Zygote and AbstractFFT rules, so if anyone has more knowledge of this please chime in.

[torfjelde]

Gotcha; thanks @ToucheSir 🙏

[devmotion]

The errors are due to the CUDA fix in JuliaMath/AbstractFFTs.jl#105 and similar conversions for adjoint plans. These changes in AbstractFFTs broke compatibility of the rules with Zygote.OneElement. I don't think you can do anything in Zygote about it - unless you want to re-add the incorrect Zygote FFT rules. It has to be fixed downstream (there are open issues and PRs).

[ToucheSir]

JuliaRegistries/General#89414

[torfjelde]

Awesome; thanks!