Overload NNlib.within_gradient

[mcabbott]

Uses FluxML/NNlib.jl#434

Does not solve #133 but the error is different:

julia> using Flux, Zygote, Tracker

julia> bn = BatchNorm(2); x = rand(Float32, 2, 3);

julia> sum(deepcopy(bn)(x))
2.1757803f0

julia> bn.active === nothing
true

julia> bn.active = false
false

julia> _bn = deepcopy(bn); Zygote.withgradient(sum∘_bn, x)[1]
2.1757803f0

julia> _bn = deepcopy(bn); Tracker.withgradient(sum∘_bn, x)[1]
2.1757803f0

julia> bn.active = true
true

julia> _bn = deepcopy(bn); Zygote.withgradient(sum∘_bn, x)[1]
-2.0489097f-7

julia> _bn = deepcopy(bn); Tracker.gradient(sum∘_bn, x)[1]
ERROR: MethodError: no method matching Float32(::Tracker.TrackedReal{Float32})
...
Stacktrace:
  [1] convert(#unused#::Type{Float32}, x::Tracker.TrackedReal{Float32})
    @ Base ./number.jl:7
  [2] setindex!
    @ ./array.jl:971 [inlined]
...
  [9] materialize!
    @ ./broadcast.jl:881 [inlined]
 [10] accum!(x::Matrix{Float32}, Δ::Matrix{Tracker.TrackedReal{Float32}})
    @ Tracker ~/.julia/packages/Tracker/a9oj5/src/back.jl:45
 [11] back(x::Tracker.Tracked{Matrix{Float32}}, Δ::Matrix{Tracker.TrackedReal{Float32}}, once::Bool)
    @ Tracker ~/.julia/packages/Tracker/a9oj5/src/back.jl:48
...
 [22] back(x::Tracker.Tracked{Matrix{Tracker.TrackedReal{Float32}}}, Δ::Matrix{Tracker.TrackedReal{Float32}}, once::Bool)

Does not make the example from FluxML/Flux.jl#2122 work, but perhaps the failure is unrelated:

julia> let
       using Flux
       using Random
       Random.seed!(123)

       model = Chain(
                 Conv((3, 3), 3 => 5, pad=1, bias=false), 
                 BatchNorm(5, relu), 
                 Conv((3, 3), 5 => 3, stride=16),
               )
       image = rand(Float32, 224, 224, 3, 1);
       @show sum(model(image));

       loss(m, x) = sum(m(x))

       opt = Flux.setup(Flux.Adam(0.001f0,  (0.9f0, 0.999f0), 1.1920929f-7), model)

       val, grads = Tracker.withgradient(model) do m
           loss(m, image)
       end

       Flux.update!(opt, model, grads[1])
       @show loss(model, image);
       end;
sum(model(image)) = -0.33076355f0
ERROR: UndefRefError: access to undefined reference
Stacktrace:
  [1] getindex
    @ ./essentials.jl:14 [inlined]
  [2] conv_direct!(y::Array{Tracker.TrackedReal{Float32}, 5}, x::Array{Tracker.TrackedReal{Float32}, 5}, w::Array{Float32, 5}, cdims::DenseConvDims{3, 3, 3, 6, 3}, ::Val{(3, 3, 1)}, ::Val{3}, ::Val{(0, 0, 0, 0, 0, 0)}, ::Val{(1, 1, 1)}, ::Val{(16, 16, 1)}, fk::Val{false}; alpha::Tracker.TrackedReal{Float32}, beta::Bool)
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:111
  [3] kwcall(::NamedTuple{(:alpha, :beta), Tuple{Tracker.TrackedReal{Float32}, Bool}}, ::typeof(NNlib.conv_direct!), y::Array{Tracker.TrackedReal{Float32}, 5}, x::Array{Tracker.TrackedReal{Float32}, 5}, w::Array{Float32, 5}, cdims::DenseConvDims{3, 3, 3, 6, 3}, ::Val{(3, 3, 1)}, ::Val{3}, ::Val{(0, 0, 0, 0, 0, 0)}, ::Val{(1, 1, 1)}, ::Val{(16, 16, 1)}, fk::Val{false})
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:59
  [4] conv_direct!(y::Array{Tracker.TrackedReal{Float32}, 5}, x::Array{Tracker.TrackedReal{Float32}, 5}, w::Array{Float32, 5}, cdims::DenseConvDims{3, 3, 3, 6, 3}; alpha::Tracker.TrackedReal{Float32}, beta::Bool)
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:50
  [5] conv_direct!
    @ ~/.julia/packages/NNlib/QJIIj/src/impl/conv_direct.jl:47 [inlined]
  [6] #conv!#301
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:288 [inlined]
  [7] conv!
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:280 [inlined]
  [8] conv!(y::Array{Tracker.TrackedReal{Float32}, 4}, x::Array{Tracker.TrackedReal{Float32}, 4}, w::Array{Float32, 4}, cdims::DenseConvDims{2, 2, 2, 4, 2}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/conv.jl:145
  [9] conv!
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:140 [inlined]
 [10] conv(x::Array{Tracker.TrackedReal{Float32}, 4}, w::Array{Float32, 4}, cdims::DenseConvDims{2, 2, 2, 4, 2}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ NNlib ~/.julia/packages/NNlib/QJIIj/src/conv.jl:88
 [11] conv
    @ ~/.julia/packages/NNlib/QJIIj/src/conv.jl:83 [inlined]
 [12] #_forward#636
    @ ~/.julia/packages/Tracker/a9oj5/src/lib/array.jl:520 [inlined]
 [13] _forward
    @ ./none:0 [inlined]
 [14] #track#1
    @ ~/.julia/packages/Tracker/a9oj5/src/Tracker.jl:58 [inlined]
 [15] track
    @ ~/.julia/packages/Tracker/a9oj5/src/Tracker.jl:57 [inlined]
 [16] #conv#633
    @ ~/.julia/packages/Tracker/a9oj5/src/lib/array.jl:516 [inlined]
 [17] conv
    @ ~/.julia/packages/Tracker/a9oj5/src/lib/array.jl:516 [inlined]
 [18] (::Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}})(x::TrackedArray{…,Array{Tracker.TrackedReal{Float32}, 4}})
    @ Flux ~/.julia/dev/Flux/src/layers/conv.jl:200
 [19] macro expansion
    @ ~/.julia/dev/Flux/src/layers/basic.jl:53 [inlined]
 [20] _applychain(layers::Tuple{Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, Bool}, BatchNorm{typeof(relu), TrackedArray{…,Vector{Float32}}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}}}, x::Array{Float32, 4})
    @ Flux ~/.julia/dev/Flux/src/layers/basic.jl:53
 [21] Chain
    @ ~/.julia/dev/Flux/src/layers/basic.jl:51 [inlined]
 [22] (::var"#loss#31")(m::Chain{Tuple{Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, Bool}, BatchNorm{typeof(relu), TrackedArray{…,Vector{Float32}}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}}}}, x::Array{Float32, 4})
    @ Main ./REPL[28]:14
 [23] (::var"#30#32"{var"#loss#31", Array{Float32, 4}})(m::Chain{Tuple{Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, Bool}, BatchNorm{typeof(relu), TrackedArray{…,Vector{Float32}}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), TrackedArray{…,Array{Float32, 4}}, TrackedArray{…,Vector{Float32}}}}})
    @ Main ./REPL[28]:19
 [24] withgradient(f::Function, xs::Chain{Tuple{Conv{2, 4, typeof(identity), Array{Float32, 4}, Bool}, BatchNorm{typeof(relu), Vector{Float32}, Float32, Vector{Float32}}, Conv{2, 4, typeof(identity), Array{Float32, 4}, Vector{Float32}}}})
    @ Tracker ~/.julia/packages/Tracker/a9oj5/src/back.jl:218
 [25] top-level scope
    @ REPL[28]:18

Edit: fixed, closes #133, closes #137

[coveralls]

Pull Request Test Coverage Report for Build 3852665337

• 13 of 16 (81.25%) changed or added relevant lines in 2 files are covered.
• 1 unchanged line in 1 file lost coverage.
• Overall coverage increased (+0.09%) to 72.439%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
src/lib/real.jl	0	1	0.0%
src/lib/array.jl	13	15	86.67%

Files with Coverage Reduction	New Missed Lines	%
src/lib/array.jl	1	67.49%

Totals
Change from base Build 3010278444:	0.09%
Covered Lines:	502
Relevant Lines:	693

---

💛 - Coveralls

[ToucheSir]

For the first error, guessing it needs something like https://github.com/FluxML/Flux.jl/blob/master/src/layers/normalise.jl#L278-L280.

For the second, this is FluxML/NNlib.jl#405. Although we probably ought to make the NNlib side more robust, the conv layer and conv functions shouldn't be seeing Array{<:TrackedReal} in the first place!

[mcabbott]

For the first, I think that after f098dd9 the error is not from within _track_stats!.

For the second, indeed the error is exactly that, thanks. But why an array of tracked numbers is made I don't know. A smaller example is this:

julia> Flux.trainmode!(BatchNorm(3))(param(rand(2,3,1)))
Tracked 2×3×1 Array{Tracker.TrackedReal{Float64}, 3}:
[:, :, 1] =
 -0.994696  -0.999959   0.912711
  0.994696   0.999959  -0.912711

julia> Flux.testmode!(BatchNorm(3))(param(rand(2,3,1)))
Tracked 2×3×1 Array{Float64, 3}:
[:, :, 1] =
 0.664405  0.684316  0.321995
 0.834291  0.463685  0.370812

... which is #137.

[ToucheSir]

Only remaining test failures on nightly are the known kron ones.