Update ChainRules for Zero -> ZeroTangent, and support a.b[constindex].c[index], fixes #278

Author: chriselrod

Project.toml

@@ -5,6 +5,7 @@ version = "0.12.29"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -20,6 +21,7 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
 ArrayInterface = "3.1.9"
+ChainRulesCore = "0.10"
 DocStringExtensions = "0.8"
 IfElse = "0.1"
 OffsetArrays = "1.4.1"
@@ -30,5 +32,5 @@ Static = "0.2"
 StrideArraysCore = "0.1.11"
 ThreadingUtilities = "0.4.2"
 UnPack = "1"
-VectorizationBase = "0.20.4"
+VectorizationBase = "0.20.4,0.21"
 julia = "1.5"

src/condense_loopset.jl

@@ -7,10 +7,9 @@ function _append_fields!(t::Expr, body::Expr, sym::Symbol, ::Type{T}) where {T}
   gf = GlobalRef(Core,:getfield)
   for f ∈ 1:fieldcount(T)
     TF = fieldtype(T, f)
+    Base.issingletontype(TF) && continue
     gfcall = Expr(:call, gf, sym, f)
-    if Base.issingletontype(TF)
-      nothing
-    elseif fieldcount(TF) ≡ 0
+    if fieldcount(TF) ≡ 0
       push!(t.args, gfcall)
     elseif TF <: DataType
       push!(t.args, Expr(:call, Expr(:curly, lv(:StaticType), gfcall)))
@@ -23,7 +22,6 @@ function _append_fields!(t::Expr, body::Expr, sym::Symbol, ::Type{T}) where {T}
   return nothing
 end
 @generated function flatten_to_tuple(r::T) where {T}
-  numfields = fieldcount(T)
   body = Expr(:block, Expr(:meta,:inline))
   t = Expr(:tuple)
   if Base.issingletontype(T)
@@ -39,10 +37,9 @@ end
   body
 end
 function rebuild_fields(offset::Int, ::Type{T}) where {T}
-  numfields = fieldcount(T)
   gf = GlobalRef(Core,:getfield)
   call = (T <: Tuple) ? Expr(:tuple) : Expr(:new, T)
-  for f ∈ 1:numfields
+  for f ∈ 1:fieldcount(T)
     TF = fieldtype(T, f)
     if Base.issingletontype(TF)
       push!(call.args, TF.instance)

src/parse/memory_ops_common.jl

@@ -1,5 +1,10 @@
-dottosym(x) = x
-dottosym(x::Expr) = Symbol(dottosym(x.args[1]), "###extractarray###", x.args[2].value)
+dottosym(x)::Symbol = x
+
+function dottosym(x::Expr)::Symbol
+  s1 = dottosym(x.args[1])
+  xa2 = x.args[2]
+  xa2 isa QuoteNode ? Symbol(s1, "###extractarray###", xa2.value) : Symbol(s1, "###extractarray###", xa2)
+end
 function extract_array_symbol_from_ref!(ls::LoopSet, ex::Expr, offset1::Int)::Symbol
     ar = ex.args[1 + offset1]
     if isa(ar, Symbol)

src/simdfunctionals/vmap_grad_rrule.jl

@@ -4,14 +4,14 @@ import .ChainRulesCore
 function ChainRulesCore.rrule(::typeof(tanh_fast), x)
     t = tanh_fast(x)
     ∂ = let t = t
-        y -> (ChainRulesCore.Zero(), mul_fast(vfnmadd_fast(t, t, one(t)), y))
+        y -> (ChainRulesCore.ZeroTangent(), mul_fast(vfnmadd_fast(t, t, one(t)), y))
     end
     t, ∂
 end
 function ChainRulesCore.rrule(::typeof(sigmoid_fast), x)
     s = sigmoid_fast(x)
     ∂ = let s = s
-        y -> (ChainRulesCore.Zero(), mul_fast(vfnmadd_fast(s, s, s), y))
+        y -> (ChainRulesCore.ZeroTangent(), mul_fast(vfnmadd_fast(s, s, s), y))
     end
     s, ∂
 end
@@ -20,7 +20,7 @@ function ChainRulesCore.rrule(::typeof(relu), v)
     cmp = v < z
     r = ifelse(cmp, z, v)
     ∂ = let cmp = cmp
-        y -> (ChainRulesCore.Zero(), ifelse(cmp, zero(y), y))
+        y -> (ChainRulesCore.ZeroTangent(), ifelse(cmp, zero(y), y))
     end
     r, ∂
 end
@@ -64,7 +64,7 @@ end
 @generated function (b::SIMDMapBack{K,T})(Δ::A) where {K,T,A}
     preloop = Expr(:block, :(jacs = b.jacs))
     loop_body = Expr(:block, :(Δᵢ = Δ[i]))
-    ret = Expr(:tuple, ChainRulesCore.Zero(), ChainRulesCore.Zero())
+    ret = Expr(:tuple, ChainRulesCore.ZeroTangent(), ChainRulesCore.ZeroTangent())
     for k ∈1:K
         jₖ = Symbol(:j_, k)
         push!(preloop.args, :($jₖ = jacs[$k]))

test/miscellaneous.jl

@@ -1069,15 +1069,15 @@ end
         @test grad!(zeros(5), ones(5), ones(3)) ≈ grad_avx!(zeros(5), ones(5), ones(3)) ≈ grad_avx_base!(zeros(5), ones(5), ones(3)) ≈ grad_avx_eval!(zeros(5), ones(5), ones(3))
 
         nta = rand(2)
-        namedtuple = (a = copy(nta), b = 10.0)
+        namedtuple = (a = (1,copy(nta)), b = 10.0)
         @turbo for i in 1:2
-            namedtuple.a[i] += namedtuple.b
+            namedtuple.a[2][i] += namedtuple.b
         end
-      @test namedtuple.a == nta .+ 10
+        @test namedtuple.a[2].c == nta .+ 10
 
-      let A = rand(T, 20, 30); B = rand(T, 20, 30); C = rand(T, 20, 30, 30);
-        @test threemulaccum_base(A,B,C) ≈ threemulaccum_lv(A,B,C)
-      end
+        let A = rand(T, 20, 30); B = rand(T, 20, 30); C = rand(T, 20, 30, 30);
+            @test threemulaccum_base(A,B,C) ≈ threemulaccum_lv(A,B,C)
+        end
     end
     for T ∈ [Int16, Int32, Int64]
         n = 8sizeof(T) - 1