Accumulate NamedTuple + Tangent

[mcabbott]

This is a hack to fix problems like this:

g_kw(;x=1.0) = sin(x)
f_kw(x) = g_kw(;x)
@test bwd(f_kw)(1.0) == bwd(sin)(1.0)  broken=true
#=
MethodError: no method matching +(::Tangent{var"#g_kw#47"{var"#g_kw#11#48"}, NamedTuple{(Symbol("#g_kw#11"),), Tuple{ZeroTangent}}}, ::Tangent{Diffractor.KwFunc{var"#g_kw#47"{var"#g_kw#11#48"}, var"#g_kw#47##kw"}, NamedTuple{(:kwf,), Tuple{ZeroTangent}}})
...
  [2] elementwise_add(a::NamedTuple{(:contents,), Tuple{Tangent{var"#g_kw#47"{var"#g_kw#11#48"}, NamedTuple{(Symbol("#g_kw#11"),), Tuple{ZeroTangent}}}}}, b::NamedTuple{(:contents,), Tuple{Tangent{Diffractor.KwFunc{var"#g_kw#47"{var"#g_kw#11#48"}, var"#g_kw#47##kw"}, NamedTuple{(:kwf,), Tuple{ZeroTangent}}}}})
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/ctmSK/src/tangent_types/tangent.jl:287
  [3] +(a::Tangent{Core.Box, NamedTuple{(:contents,), Tuple{Tangent{var"#g_kw#47"{var"#g_kw#11#48"}, NamedTuple{(Symbol("#g_kw#11"),), Tuple{ZeroTangent}}}}}}, b::Tangent{Core.Box, NamedTuple{(:contents,), Tuple{Tangent{Diffractor.KwFunc{var"#g_kw#47"{var"#g_kw#11#48"}, var"#g_kw#47##kw"}, NamedTuple{(:kwf,), Tuple{ZeroTangent}}}}}})
    @ ChainRulesCore ~/.julia/packages/ChainRulesCore/ctmSK/src/tangent_arithmetic.jl:130
=#

Diffractor sometimes produces ChainRulesCore's Tangent types, and sometimes produces plain NamedTuples instead. If these happen for structural gradients of the same object, then you will get an error.

Ideally it should probably make up its mind. Maybe Tangents are a pain to deal with at higher order?

[codecov-commenter]

Codecov Report

Merging #88 (c9f2a90) into main (9a8a788) will decrease coverage by 2.19%.
The diff coverage is 72.72%.

@@            Coverage Diff             @@
##             main      #88      +/-   ##
==========================================
- Coverage   53.89%   51.69%   -2.20%     
==========================================
  Files          21       21              
  Lines        2171     2124      -47     
==========================================
- Hits         1170     1098      -72     
- Misses       1001     1026      +25     

Impacted Files	Coverage Δ
src/runtime.jl	78.94% <70.00%> (+20.61%)	⬆️
src/extra_rules.jl	35.45% <100.00%> (-0.33%)	⬇️
src/stage1/recurse.jl	91.48% <0.00%> (-5.29%)	⬇️
src/stage1/forward.jl	69.38% <0.00%> (-4.77%)	⬇️
src/stage1/hacks.jl	2.12% <0.00%> (-4.26%)	⬇️
src/tangent.jl	32.97% <0.00%> (-1.07%)	⬇️
src/stage1/generated.jl	72.72% <0.00%> (-0.58%)	⬇️
src/jet.jl	39.16% <0.00%> (-0.51%)	⬇️
src/stage1/recurse_fwd.jl	94.11% <0.00%> (-0.17%)	⬇️
... and 3 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[mcabbott]

I get some errors from this adding tuples of mis-matched lengths, which I haven't tracked down. Not touched in this PR.

What I don't know here is whether accum is only ever called on a single tangent, or also on a tuple of a function's arguments. IIRC there is no distinction in Zygote, but in CR the former can be Array or Tangent or NoTangent, but the latter can only ever be a Tuple (of the right length).

[mcabbott]

This is JuliaDiff/ChainRulesCore.jl#581 now

[oscardssmith]

I don't think this is the right long term solution, but it seems right for now.

[mcabbott]

Yes for sure it's a hack. But it doesn't seem like it will hide serious bugs, and having tests that work for now is good.

Last commit disables tests on 1.8, since judging by "┌ Warning: ir verification broken. Either use 1.9 or 1.7" there isn't an intention to make everything work there. (The actual failures are inference tests.)

Edit: and now nightly has some unexpected passes, inference tests... Maybe we need to split CI into two files, to allow failures on 1.8 + nightly while returning a useful overall pass/fail?

[ToucheSir]

On the topic of longer-term solutions, why is Box appearing here at all? kwsorters are usually type stable and Diffractor doesn't appear to change their branching in ways which would disrupt that. Diffractor.KwFunc also doesn't make use of any inner closures.

[mcabbott]

I don't know. Where should we put things like this, where a failing test has uncovered bad behaviour, but possibly no longer bad enough to make the test fail?

(Fun Zygote example where re-using a variable name caused a 5x slowdown for ages, later became a bug: Z#1290)