Rules for eachslice (#561)

* rules for eachslice

* test only 1.6+

* move methods, use projection, allow 2nd derivatives

* fixup

* avoid limitations of ComposedFunction on 1.6

* fixup rebase, and bump version

* change to at-assert

* Apply 5 suggestions

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Co-authored-by: Lyndon White <oxinabox@ucc.asn.au>

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.24"
+version = "1.25"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/Base/indexing.jl

@@ -48,4 +48,68 @@ function frule((_, ẋ, v̇), ::typeof(setindex!), x::AbstractArray, v, inds...)
 end
 
 
+#####
+##### `eachslice` and friends
+#####
+
+function rrule(::typeof(eachrow), x::AbstractVecOrMat)
+    allrows(dy) = (NoTangent(), ∇eachslice(unthunk(dy), x, Val(1)))
+    return collect(eachrow(x)), allrows
+end
+
+function rrule(::typeof(eachcol), x::AbstractVecOrMat)
+    allcols(dy) = (NoTangent(), ∇eachslice(unthunk(dy), x, Val(2)))
+    return collect(eachcol(x)), allcols
+end
+
+function rrule(::typeof(eachslice), x::AbstractArray; dims)
+    y = collect(eachslice(x; dims=dims))
+    @assert length(dims) == 1 """That's amazing, after many years JuliaLang/julia#32310
+        actually landed. Sadly, the gradient rule for `eachslice` is unable to handle this
+        case right now, please make an issue at https://github.com/JuliaDiff/ChainRules.jl"""
+    dim = only(dims)
+    allslices(dy) = (NoTangent(), ∇eachslice(unthunk(dy), x, Val(dim)))
+    return y, allslices
+end
+
+# Using Val(dim) here is worth a factor of 2 in this, on Julia 1.8-
+# @btime rrule(eachcol, $([1 2; 3 4]))[2]($([[10, 20], [30, 40]]))
+function ∇eachslice(dys_raw, x::AbstractArray, vd::Val{dim}) where {dim}
+    dys = unthunk(dys_raw)
+    i1 = findfirst(dy -> dy isa AbstractArray, dys)
+    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))
+    # 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)
+        slice = selectdim(dx, dim, i)
+        if dys[i] isa AbstractZero
+            _zero_fill!(slice)  # Avoids this: copyto!([1,2,3], ZeroTangent()) == [0,2,3]
+        else
+            copyto!(slice, dys[i])
+        end
+    end
+    return ProjectTo(x)(dx)
+end
+
+_zero_fill!(dx::AbstractArray{<:Number}) = fill!(dx, zero(eltype(dx)))
+_zero_fill!(dx::AbstractArray) = map!(zero, dx, dx)
+
+function rrule(::typeof(∇eachslice), dys, x, vd::Val)
+    function ∇∇eachslice(dz_raw)
+        dz = unthunk(dz_raw)
+        # eachslice(dz; dims=_val(vd)) does not make @code_warntype happy
+        iter = vd == Val(1) ? eachrow(dz) : vd == Val(2) ? eachcol(dz) : eachslice(dz; dims=_val(vd))
+        return (NoTangent(), collect(iter), NoTangent(), NoTangent())
+    end
+    return ∇eachslice(dys, x, vd), ∇∇eachslice
+end
 
+# These rules help with testing, and won't hurt:
+# They are correct as we always `collect` the primal result as we need that
+# information for the reverse pass
+ChainRules.rrule(::typeof(collect∘eachrow), x) = rrule(eachrow, x)
+ChainRules.rrule(::typeof(collect∘eachcol), x) = rrule(eachcol, x)
+ChainRules.rrule(::typeof(collect∘eachslice), x; dims) = rrule(eachslice, x; dims=dims)

test/rulesets/Base/indexing.jl

@@ -68,3 +68,34 @@ end
     test_frule(setindex!, rand(3, 4), rand(), 1, 2)
     test_frule(setindex!, rand(3, 4), [1,10,100.0], :, 3)
 end
+
+@testset "eachslice" begin
+    # Testing eachrow not collect∘eachrow leads to errors, e.g.
+    # test_rrule: eachrow on Vector{Float64}: Error During Test at /Users/me/.julia/packages/ChainRulesTestUtils/8dFTY/src/testers.jl:195
+    #   Got exception outside of a @test
+    #   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))
+
+    test_rrule(collect∘eachcol, rand(3, 4))
+    @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,)))
+    end
+
+    # Make sure pulling back an array that mixes some AbstractZeros in works right
+    _, back = rrule(eachcol, rand(3, 4))
+    @test back([1:3, ZeroTangent(), 7:9, NoTangent()]) == (NoTangent(), [1 0 7 0; 2 0 8 0; 3 0 9 0])
+    @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])
+
+    # 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)
+end