Fix computing of array element type for ∇eachslice

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.70.0"
+version = "1.71.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/rulesets/Base/indexing.jl

@@ -267,7 +267,7 @@ function ∇eachslice(dys_raw, x::AbstractArray, vd::Val{dim}) where {dim}
     if i1 === nothing  # all slices are Zero!
         return _zero_fill!(similar(x, float(eltype(x)), axes(x)))
     end
-    T = promote_type(eltype(dys[i1]), eltype(x))
+    T = Base.promote_eltype(dys...)
     # The whole point of this gradient is that we can allocate one `dx` array:
     dx = similar(x, T, axes(x))
     for i in axes(x, dim)
@@ -282,8 +282,12 @@ function ∇eachslice(dys_raw, x::AbstractArray, vd::Val{dim}) where {dim}
 end
 ∇eachslice(dys::AbstractZero, x::AbstractArray, vd::Val{dim}) where {dim} = dys
 
-_zero_fill!(dx::AbstractArray{<:Number}) = fill!(dx, zero(eltype(dx)))
-_zero_fill!(dx::AbstractArray) = map!(zero, dx, dx)
+_zero_fill!(dx::AbstractArray) = fill!(dx, zero(eltype(dx)))
+
+# Belong in ChainRulesCore
+Base.promote_rule(T::Type{<:Number}, S::Type{<:AbstractZero}) = T
+Base.promote_rule(T::Type{<:AbstractZero}, S::Type{<:Number}) = S
+Base.eltype(::Type{NoTangent}) = NoTangent
 
 function rrule(::typeof(∇eachslice), dys, x, vd::Val)
     function ∇∇eachslice(dz_raw)

test/rulesets/Base/indexing.jl

@@ -217,16 +217,24 @@ end
     #   DimensionMismatch("second dimension of A, 6, does not match length of x, 5")
     # Probably similar to https://github.com/JuliaDiff/ChainRulesTestUtils.jl/issues/234 (about Broadcasted not Generator)
 
-    test_rrule(collect∘eachrow, rand(5))
-    test_rrule(collect∘eachrow, rand(3, 4))
+    # Inference on 1.6 sometimes fails, so don't enforce there.
+    test_rrule(collect ∘ eachrow, rand(5); check_inferred=(VERSION >= v"1.7"))
+    test_rrule(collect ∘ eachrow, rand(3, 4); check_inferred=(VERSION >= v"1.7"))
 
-    test_rrule(collect∘eachcol, rand(3, 4))
-    @test_skip test_rrule(collect∘eachcol, Diagonal(rand(5)))  # works locally!
+    test_rrule(collect ∘ eachcol, rand(3, 4); check_inferred=(VERSION >= v"1.7"))
+    @test_skip test_rrule(collect ∘ eachcol, Diagonal(rand(5)))  # works locally!
 
     if VERSION >= v"1.7"
         # On 1.6, ComposedFunction doesn't take keywords. Only affects this testing strategy, not real use.
-        test_rrule(collect∘eachslice, rand(3, 4, 5); fkwargs = (; dims = 3))
-        test_rrule(collect∘eachslice, rand(3, 4, 5); fkwargs = (; dims = (2,)))
+        test_rrule(collect ∘ eachslice, rand(3, 4, 5); fkwargs=(; dims=3))
+        test_rrule(collect ∘ eachslice, rand(3, 4, 5); fkwargs=(; dims=(2,)))
+
+        test_rrule(
+            collect ∘ eachslice,
+            FooTwoField.(rand(3, 4, 5), rand(3, 4, 5));
+            check_inferred=false,
+            fkwargs=(; dims=3),
+        )
     end
 
     # Make sure pulling back an array that mixes some AbstractZeros in works right
@@ -235,8 +243,16 @@ end
     @test back([1:3, ZeroTangent(), 7:9, NoTangent()])[2] isa Matrix{Float64}
     @test back([ZeroTangent(), ZeroTangent(), NoTangent(), NoTangent()]) == (NoTangent(), [0 0 0 0; 0 0 0 0; 0 0 0 0])
 
+    _, back = ChainRules.rrule(
+        eachslice, FooTwoField.(rand(2, 3, 2), rand(2, 3, 2)); dims=3
+    )
+    @test back([fill(Tangent{Any}(; x=0.0, y=1.0), 2, 3), fill(ZeroTangent(), 2, 3)]) == (
+        NoTangent(),
+        cat(fill(Tangent{Any}(; x=0.0, y=1.0), 2, 3), fill(ZeroTangent(), 2, 3); dims=3),
+    )
+
     # Second derivative rule
     test_rrule(ChainRules.∇eachslice, [rand(4) for _ in 1:3], rand(3, 4), Val(1))
     test_rrule(ChainRules.∇eachslice, [rand(3) for _ in 1:4], rand(3, 4), Val(2))
-    test_rrule(ChainRules.∇eachslice, [rand(2, 3) for _ in 1:4], rand(2, 3, 4), Val(3), check_inferred=false)
+    test_rrule(ChainRules.∇eachslice, [rand(2, 3) for _ in 1:4], rand(2, 3, 4), Val(3); check_inferred=(VERSION >= v"1.7"))
 end