WIP: use ChainRules in Forward #752

willtebbutt · 2020-08-01T14:21:29Z

Addresses #712.

TODO:

Get nested AD working
Ensure that kwargs work correctly
"Unwrap" ChainRules types into Zygote type for compatibility
Remove all excess rules
Tidy up

A separate PR will be needed to move rules that are currently in Zygote.Forward to ChainRules. Unlike the reverse-mode rules there aren't a particularly large number of them, so it shouldn't be a massive job.

willtebbutt

This is an implementation in line with the way that the reverse-mode one is done. Once JuliaDiff/ChainRulesCore.jl#182 is complete, we could also just utilise that, and I think the implementation will probably be a bit simpler.

willtebbutt · 2020-08-01T14:23:13Z

src/forward/compiler.jl

+ return :(__pushforward(dargs, f, args...))
+ end
+
+ # g = try _lookup_grad(T) catch e e end


I'm really not sure what I'm doing with the edges stuff here -- this is currently copied from the reverse-mode definition. @oxinabox or @DhairyaLGandhi could you comment on what sort of thing I'm going to need to do here? In particular, I'm not sure what _lookup_grad is doing, so I'm not sure how to reason about adapting the reverse-mode implementation.

willtebbutt · 2020-08-01T14:25:35Z

test/forward/forward.jl

+#  end
+#  end
+#  x
+# end == 0

 @test D(x -> abs(x+2im), 1) == gradient(x -> abs(x+2im), 1)[1]


This doesn't currently pass. I'm pretty sure that the Forward implementation is doing the correct thing, and gradient is wrong. D (correctly) returns a Real as the derivative w.r.t x, while gradient returns a Complex. @sethaxen do you agree that Zygote is doing the wrong thing here?

ChainRulesCore has an issue open about this: JuliaDiff/ChainRulesCore.jl#176. It's also related to JuliaDiff/ChainRules.jl#232. The issue is that our rules for addition between real and complex (+(x, 2im)) will pull back a complex adjoint to x.
So we'd expect these to do the same thing, but

julia> gradient(x -> abs(x+2im), 1)[1] 0.4472135954999579 + 0.8944271909999159im julia> gradient(x -> abs(complex(x, 2)), 1)[1] 0.4472135954999579

The first issue above links to some Zygote discussions. Essentially, Zygote treats all reals as actually complex, so from that perspective it is doing the right thing. I just don't think that's the right way to go.

Ahh interesting. It's interesting that Zygote.Forward doesn't naturally do the same thing as Zygote.

I guess that the question is what the right way forward for us is with Forward. Mike previously had some extra code for abs to make it do the Zygote-y thing.

I wonder whether it's acceptable to utilise a different convention in Forward than what Zygote currently does...

@DhairyaLGandhi what are your thoughts on the way that Zygote currently handles complex numbers / have you had a chance to review the various ChainRules discussions about complex numbers?

test/forward/forward.jl

willtebbutt · 2020-08-01T15:27:22Z

Most recent commit fixed one bug, but re-introduced some tests that expose us to having to differentiate through broadcast, which isn't something that Zygote.Forward currently handles. I'll need to give that some thought...

oxinabox · 2020-08-03T13:05:15Z

src/forward/compiler.jl

+
+"""
+is_kwfunc(::Vararg{Any}) = false
+is_kwfunc(k, ::Type{<:NamedTuple}, f, args...) = k===Core.kwftype(f)


this code exists already for reverse mode we should just import it from there?

Yes. I'll need to refactor to avoid all of this duplication. I just copied + pasted to start with because I was tired haha

oxinabox · 2020-08-03T13:07:29Z

src/forward/compiler.jl

+end
+
+"""
+ chain_frule(f, args...)


this doc is wrong but the code is right.
dargs includes dself, darg1, darg2 ... and args includes foo, darg1, darg2

oxinabox · 2020-08-03T13:09:21Z

src/forward/compiler.jl

+
+ chain_frule_f = iskw ? :chain_frule_kw : :chain_frule
+ if hascr
+ return :($chain_frule_f(dargs, f, args...))


where are my kwargs gone?

Just not implemented yet, see check-boxes at the top :)

src/forward/compiler.jl

oxinabox · 2020-08-03T13:11:37Z

src/forward/compiler.jl

+function has_chain_frule(T)
+ m = meta(Tuple{typeof(frule), T.parameters...})
+
+ if m.method !== chainrules_frule_fallback


I wonder if we should just abstract this to be has_chainrule(kind, T)

Suggested change

function has_chain_frule(T)

m = meta(Tuple{typeof(frule), T.parameters...})

if m.method !== chainrules_frule_fallback

function has_chain_rule(kind, T)

m = meta(Tuple{typeof(kind), T.parameters...})

if m.method !== chainrules_rule_fallback(kind)

oxinabox · 2020-08-03T13:12:29Z

test/forward/forward.jl

+foo(x) = x
+
+# This intentionally has the wrong definition so that we can detect if Zygote is using it.
+ChainRulesCore.frule((_, dx), ::typeof(foo), x) = x, 2 * dx


the tests in the reverse mode chainrules tests set a variable that can late be read to see if it was hit.
This should be tested the same way

Oh nice, I'll do that.

oxinabox · 2020-08-03T13:13:20Z

src/lib/lib.jl

@@ -76,6 +76,10 @@ function global_set(ref, val)
 ref.mod, ref.name, val)
 end

+function ChainRules.frule(dargs, ::typeof(Zygote.global_set), ref, val::Nothing)
+ return Zygote.global_set(ref, val), nothing


Suggested change

return Zygote.global_set(ref, val), nothing

return Zygote.global_set(ref, val), DoesNotExist()

oxinabox · 2020-08-03T13:14:42Z

src/forward/compiler.jl

+Returns a the (primal) value of `f(args...)` and tangent, by invoking
+`ChainRules.frule(f, args...)`.
+"""
+@noinline chain_frule(dargs, args...) = frule(dargs, args...)


we need a translation layer replacing Composites with NamedTuples and AbstractZeros with nothing
and visa versa

Like we have for reverse mode.
Can probably just abstract and call that

Yup, on the TODO list :)

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

willtebbutt · 2020-08-03T14:15:31Z

Thanks for the review @oxinabox . The main thing I'm still unsure about is what _lookup_grad was doing in the reverse mode implementation. (See my first comment)

oxinabox · 2020-08-05T19:36:23Z

AFAICT _lookup_grad in reverse, dispite the name, does not look up the gradient,
it performs the core of the source code tranform to do AD. it is defined in emit.jl
Roughly speaking the matching function in forward is dual which is also terribly named

willtebbutt · 2020-09-23T16:39:00Z

Closing because I've definitely not got time to work on this at the minute, and don't anticipate doing so for a while. If someone else wants to have a go at this, they should feel free to do so.

Refactor compiler + remove DiffRules

a56b41b

willtebbutt commented Aug 1, 2020

View reviewed changes

willtebbutt requested review from DhairyaLGandhi and oxinabox August 1, 2020 14:31

willtebbutt added 2 commits August 1, 2020 16:19

Tidy up docstrings slightly

5d27035

Fix a higher-order bug

498ccb9

oxinabox reviewed Aug 3, 2020

View reviewed changes

Update src/forward/compiler.jl

786ec20

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

willtebbutt closed this Sep 23, 2020

willtebbutt deleted the wct/forward-cr branch January 16, 2021 12:44

ToucheSir mentioned this pull request Dec 11, 2021

Broadcast gradient error #1132

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

WIP: use ChainRules in Forward #752

WIP: use ChainRules in Forward #752

willtebbutt commented Aug 1, 2020 •

edited

Loading

willtebbutt left a comment

willtebbutt Aug 1, 2020

willtebbutt Aug 1, 2020

sethaxen Aug 1, 2020

willtebbutt Aug 2, 2020

willtebbutt commented Aug 1, 2020

oxinabox Aug 3, 2020

willtebbutt Aug 3, 2020

oxinabox Aug 3, 2020

oxinabox Aug 3, 2020

willtebbutt Aug 3, 2020

oxinabox Aug 3, 2020

oxinabox Aug 3, 2020

willtebbutt Aug 3, 2020

oxinabox Aug 3, 2020

oxinabox Aug 3, 2020

willtebbutt Aug 3, 2020

willtebbutt commented Aug 3, 2020

oxinabox commented Aug 5, 2020

willtebbutt commented Sep 23, 2020

	return Zygote.global_set(ref, val), nothing
	return Zygote.global_set(ref, val), DoesNotExist()

WIP: use ChainRules in Forward #752

WIP: use ChainRules in Forward #752

Conversation

willtebbutt commented Aug 1, 2020 • edited Loading

willtebbutt left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

willtebbutt commented Aug 1, 2020

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

willtebbutt commented Aug 3, 2020

oxinabox commented Aug 5, 2020

willtebbutt commented Sep 23, 2020

willtebbutt commented Aug 1, 2020 •

edited

Loading