Open sethaxen opened 4 years ago
OK - I'm not at all familiar with Zygote. Would supporting it mean (a) adding Zygote as a dependency and (b) adding some new methods to some Zygote functions (like rrules
)?
OK - I'm not at all familiar with Zygote. Would supporting it mean (a) adding Zygote as a dependency and (b) adding some new methods to some Zygote functions (like
rrules
)?
Nope! It would only require adding ChainRulesCore as a dependency, which has a single dependency that in turn has no dependencies, so it's extremely light. The rrule
function is in ChainRulesCore, so it would only involve overriding rrule
for the constructors.
@oxinabox gave a nice talk on ChainRules at JuliaCon this year: https://www.youtube.com/watch?v=B4NfkkkJ7rs
Yeah OK, seems quite reasonable to me. (Pull requests are always welcome :wink:)
Supporting reverse-mode autodiff with Zygote requires two things:
I think the latter is the main thing missing for Zygote support. Supporting it would require adding custom
rrule
s for constructors using ChainRulesCore, which has essentially no dependencies.Here's an example that fails:
In this case mathematically the missing pullback is for the exponential map
exp: so(3) → SO(3)
.