search

FluxML / Zygote.jl

21st century AD
https://fluxml.ai/Zygote.jl/
Other
1.48k stars 210 forks source link

Extend kron support #1458

Closed willtebbutt closed 1 year ago

willtebbutt commented 1 year ago

The current kron rules do not handle kron(::AbstractMatrix, ::AbstractVector) or kron(::AbstractVector, ::AbstractMatrix). This PR adds support for both. Fortunately, the implementation strategy is precisely the same as the one currently taken.

A method of _pullback for kron is removed in favour of a single method of _pullback which can handle all four combinations of argument types.

Fixes https://github.com/JuliaGaussianProcesses/TemporalGPs.jl/pull/115

PR Checklist

willtebbutt commented 1 year ago

Has it caused you problems previously?

simsurace commented 1 year ago

I can't wait for us to be rid of this rule

Maybe in the not-so-distant future: https://github.com/JuliaDiff/ChainRules.jl/pull/741

ToucheSir commented 1 year ago

Has it caused you problems previously?

Not so much. It was just frustrating to have CI broken for so long (in Zygote and Tracker) because of changes in Base and to have to patch the issue in the AD itself instead of in ChainRules. I'm glad to see progress is being made on a proper rule :)