Interface documentation about zygote.jl HOT 4 CLOSED

Just added fairly comprehensive interface documentation to the docs. Let me know if there's anything unclear in that.

from zygote.jl.

No other docs yet, but I can try to clarify these and figure out what needs documenting explicitly.

In general we specify the args we want using closures, e.g.

f(a, b) = a*b
gradient(f, 2, 3) # (3, 2)
gradient(x -> f(2, x), 3) # (2,)

This is one of those things that makes total sense but is a bit non-obvious at first, so it'd be worth having explicit docs on this.

I'm a bit confused by the other two questions; the gradients are right there, and running without gradients is just f(2, 3) as usual. Maybe you're referring to the Params API which is a bit different.

It might be useful to see how I'd simulate Knet/AutoGrad's original API with Zygote:

grad(f, arg = 1) = (x...) -> gradient(f, x...)[arg]

loss(θ, x) = sum(θ[:W] * x .+ θ[:b])

∇loss = grad(loss)

θ = Dict(:W => randn(5,5), :b => randn(5))
# θ = (W = randn(5,5), b = randn(5)) # NamedTuples work too

∇loss(θ, rand(5)) # Dict(:W => ..., :b => ...)

from zygote.jl.

from zygote.jl.

All fair points. I've mostly been focusing on compiler crashes up to now, but things are getting more usable so it's time to add these kinds of docstrings.

Yes, you can always take the gradient of a structure; things like dropping keys are technically a library-level issue but we will drop keys of dicts right now. We also return nothing in the case that a value is not used at all; I'm not completely set on that but up to now it seems like a reasonable way forward.

I am wondering if we can write a 5 line mini-Zygote example in Julia

Yeah this is a good idea. It's not easy to write down the code transformation, but it is pretty easy to simulate what Zygote is doing by hand, so it'd be good to have a notebook or some docs on that.

from zygote.jl.