using Zygote

[MikeInnes]

Otherwise known as "break all the things". This will be a huge change so I'm beginning to prepare now, even though Zygote is still a couple of months off from being really ready. Do not try this at home (yet) – this branch is eventually aimed at beta testers, but isn't even ready for that yet.

The idea is to break as little code as possible, which means supporting the current Params API; but I also want to start prototyping the nicer things discussed in #628 and other issues.

Blocking issues:

• Get the tests passing.
• Check tests on GPU.
• Rewrite all the docs.
• Cache invalidation (world age validation for certain generated functions JuliaLabs/Cassette.jl#6).
• Moving over adjoints (Move over adjoints from Flux Zygote.jl#81).
• General Zygote robustness.

Nice to have:

• Robust nested AD (may not be a blocker if one can still use Tracker with Flux).
• Zygote support for modules / globals as discussed in Gradient Interface Design #628, along with New New Optimisers #637.
• Better train/test mode as in Train/test mode #643.

If you're the kind of person who ignores triangular road signs, you can try this with

]add Flux#zygote Zygote#master

[MikeInnes]

I initially had started doing #666 and #637 on this branch, but that's turning into a big project, so I've stripped the commits for now (preserved on mji/step). As excited as I am to get rid of Params I think the right move is to shelve that for now and focus on the core AD issues.

[staticfloat]

In my own experiments trying to use Zygote with Flux (I just do model = mapleaves(Flux.data, model) first, then define my own Zygote-based update step:

zyg_update!(opt, model, updates::Nothing) = nothing
function zyg_update!(opt, model::AbstractArray, updates::AbstractArray)
    # Sub off to Flux's ADAM optimizer
    Δ = Flux.Optimise.update!(opt, model, updates)
    return model .-= Δ
end

function zyg_update!(opt, model, updates)
    if nfields(model) == 0
        return model
    end

    for field_idx in 1:nfields(model)
        zyg_update!(opt, getfield(model, field_idx), getfield(updates, field_idx))
    end
end

Things actually work fairly well, except BatchNorm freaks out, complaining about mutating arrays. To work around this, I am using my own BatchNorm implementation, re-architected to work with Zygote. Not sure if that is the direction you want to go with this Mike, but it worked well for us on TPU. I will note, anecdotally, that my convolution and batchnom-heavy workload (a convolutional autoencoder for large images) uses ~20% less memory with Zygote than with Tracker.

[staticfloat]

I wanted to use this with the new NNlib overhaul, so I rebased this branch on top of master; I'm not certain I did everything right, but sf/zygote_updated contains my rebased version. Mike, if you like it, you can just force-push it to #zygote and keep on working, or you can just do the rebase yourself.

[MikeInnes]

@dhairyagandhi96 Flux already exports gradient, so this may not be necessary

[DhairyaLGandhi]

Fixed

[MikeInnes]

bors r+

[MikeInnes]

Seem to be some issues with our GPU CI, so just merging.

[bors]

Build failed

• ci/gitlab/staging

[iblislin]

well, is there an alternative?