Update rrules so that we can support Yota

[avik-pal]

Note that this is still not enough. Some additional patches are still needed (mostly upstream?)

function Base.:+(t::ChainRulesCore.Tangent, nt::NamedTuple)
    return merge(ChainRulesCore.backing(t), nt)
end

function Yota.bcast_rrule(rc::Yota.YotaRuleConfig, ::typeof(Base.Broadcast.broadcasted),
                          f::typeof(sqrt), args...; kw...)
    ys, pbs = Yota.unzip(ChainRulesCore.rrule.(rc, f, args...; kw...))
    function pullback(Δ)
        if Δ isa ChainRulesCore.NoTangent || Δ isa ChainRulesCore.ZeroTangent
            return (ChainRulesCore.NoTangent(), [Δ for _ in 1:(length(pbs) + 1)]...)
        end
        Δ = ChainRulesCore.unthunk(Δ)
        dxs = broadcast((pb, Δ) -> pb(Δ), pbs, Δ) |> Yota.unzip
        dxs = [all(dx .== ChainRulesCore.NoTangent()) ? ChainRulesCore.NoTangent() : dx
               for dx in dxs]
        return (ChainRulesCore.NoTangent(), dxs...)
    end
    return ys, pullback
end

function Yota.ungetindex(x::NamedTuple, dy, I...)
    dx = map(1:length(x)) do i
        return i in I ? dy : ChainRulesCore.ZeroTangent()
    end
    return NamedTuple{keys(x)}(dx)
end

Base.zero(::Val{T}) where {T} = Val(zero(T))
Base.zero(::Nothing) = nothing
Base.zero(nt::NamedTuple) = Functors.fmap(zero, nt)

[codecov]

Codecov Report

Merging #85 (31cf81d) into main (f7ba291) will increase coverage by 5.03%.
The diff coverage is 86.79%.

@@            Coverage Diff             @@
##             main      #85      +/-   ##
==========================================
+ Coverage   71.18%   76.21%   +5.03%     
==========================================
  Files          14       14              
  Lines         989     1026      +37     
==========================================
+ Hits          704      782      +78     
+ Misses        285      244      -41     

Impacted Files	Coverage Δ
src/core.jl	100.00% <ø> (ø)
src/autodiff.jl	65.82% <79.41%> (+9.30%)	⬆️
src/layers/normalize.jl	86.77% <100.00%> (+35.49%)	⬆️
src/nnlib.jl	64.86% <100.00%> (+1.80%)	⬆️
src/layers/basic.jl	62.63% <0.00%> (+2.63%)	⬆️

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f7ba291...31cf81d. Read the comment docs.