Gradients for prod(x; dims) and cumprod(x), which take keywords & allow for zero entries

src/tracker/array.jl

@@ -213,15 +213,84 @@ Base.sum(f::Union{Function,Type},xs::TrackedArray) = sum(f.(xs))
 @grad sum(xs; dims = :) = sum(data(xs), dims = dims),
   Δ -> (zero(xs) .+ Δ, )
 
-Base.prod(xs::TrackedArray, dim) = track(prod, xs, dim)
-Base.prod(xs::TrackedArray) = track(prod, xs)
+Base.prod(xs::TrackedArray; dims=:) = track(prod, xs; dims=dims)
 Base.prod(f::Union{Function, Type}, xs::TrackedArray) = prod(f.(xs))
 
-@grad prod(xs) = prod(data(xs)), Δ -> (prod(xs) ./ xs .* Δ,)
-@grad prod(xs, dim) = prod(data(xs), dims = dim),
-  Δ -> (nobacksies(:sum,
-          reshape(.*(circshift.([reshape(data(xs), length(xs))], 1:length(xs)-1)...), size(xs)) .* Δ),
-        nothing)
+@grad function prod(xs; dims=:) 
+  p = prod(xs.data, dims=dims)
+  length(p)==1 ? _prod(xs.data, first(p), :) : _prod(xs.data, p, dims) # avoid mapslices() if not necc.
+end
+_prod(xs::Array, p, ::Colon) = p, Δ -> (nobacksies(:prod, ∇prod(xs, p, data(Δ)) ),) 
+_prod(xs::Array, p, d::Int) = p, Δ -> (nobacksies(:prod, ∇prod(xs, Val(d), p, data(Δ)) ), )
+_prod(xs::Array, p, dims) = p, Δ -> (nobacksies(:prod, mapslices(∇prod, xs; dims=dims) .* Δ), )
+_prod(xs, p, ::Colon) =  p, Δ -> ( ∇prod_all(x) .* Δ ,) # more generic fall-back
+_prod(xs, p, dim) =      p, Δ -> ( ∇prod_dim(x, p, dim) .* Δ ,)
+
+function ∇prod(x::AbstractArray, p=prod(x), Δ=1, f=identity)
+  !iszero(p) && return p ./ x .* Δ
+  ∇ = fill!(similar(x), 0)
+  z = findall(iszero, x)
+  length(z)>1 && return ∇
+  ∇[first(z)] = prod(f(xi) for (i, xi) in pairs(x) if i!=first(z)) * Δ
+  return ∇ 
+end
+
+∇prod_all(x) = reshape( .*(circshift.([reshape(x, length(x))], 1:length(x)-1)...), size(x))
+∇prod_dim(x, p, dim::Int) = .*(circshift.([x], tuple.(Iterators.repeated(0,dim-1)..., 1:size(x,dim)-1))...)
+∇prod_dim(x, p, dim) = p ./ x # error if x contains zero
+
+function ∇prod(x::AbstractArray, vald::Val{d}, p=prod(x, dims=d), Δ=fill!(similar(p),1)) where d
+  ∇ = similar(x)
+  for ii in Iterators.product(tup_insert(axes(x), Base.OneTo(1), vald)...)
+    iic = tup_insert(ii, Colon(), vald)
+    iio = tup_insert(ii, 1, vald)
+    ∇[iic...] .= ∇prod(x[iic...], p[iio...], Δ[iio...]) # like mapslices(∇prod, x,p,Δ; dims=d) if that existed
+  end
+  ∇
+end
+
+tup_insert(tup::NTuple{N}, what, ::Val{d}) where {N,d} = ntuple(i -> ifelse(i==d, what, tup[i]), Val(N))
+
+Base.cumprod(xs::TrackedArray; dims::Union{Nothing, Integer}=nothing) = track(cumprod, xs; dims=dims)
+
+@grad function cumprod(xs; dims=nothing) 
+  ndims(xs)==1 && (dims==1 || dims==(1,)) && return _cumprod(xs, nothing)
+  _cumprod(xs, dims)
+end
+_cumprod(xs, ::Nothing) = begin p = cumprod(xs.data); p, Δ -> (nobacksies(:cumprod, ∇cumprod(xs.data, p, data(Δ)) ),) end
+_cumprod(xs, d) = begin p = cumprod(xs.data, dims=d); p, Δ -> (nobacksies(:cumprod, ∇cumprod(xs.data, Val(d), p, data(Δ)) ),) end
+
+function ∇cumprod(x::AbstractVector, p=cumprod(x), Δ=fill(1, length(x)))
+  len = length(x)
+  z = something(findfirst(iszero, x), len+1)
+
+  ∇ = fill!(similar(x), 0)
+  @inbounds for i=1:z-1
+    ixi = 1/x[i]
+    for k=i:z-1
+      ∇[i] += p[k] * Δ[k] * ixi
+    end
+  end
+
+  @inbounds if z != len+1
+    pk = z==1 ? one(p[1]) : p[z-1] # will be prod(x[j] for j=1:k if j!=z)
+    ∇[z] += pk * Δ[z]
+    for k=(z+1):len
+      pk *= x[k]
+      ∇[z] += pk * Δ[k]
+    end
+  end
+  ∇
+end
+
+function ∇cumprod(x::AbstractArray, vald::Val{d}, p=cumprod(x; dims=d), Δ=fill(1, size(x))) where d
+  ∇ = similar(x)
+  for ii in Iterators.product(tup_insert(axes(x), Base.OneTo(1), vald)...)
+    iid = tup_insert(ii, Colon(), vald)
+    ∇[iid...] .= ∇cumprod(x[iid...], p[iid...], Δ[iid...])
+  end
+  ∇
+end
 
 Base.findfirst(xs::TrackedArray, args...) = findfirst(xs.data, args...)
 

test/tracker.jl

@@ -17,12 +17,37 @@ gradtest(f, dims...) = gradtest(f, rand.(Float64, dims)...)
 @test gradtest((x, W, b) -> logσ.(W*x .+ b), (5,3), (2,5), 2)
 @test gradtest((w, x) -> w'*x, randn(Float64,10, 2), randn(Float64,10))
 @test gradtest((w, x) -> w*x', randn(Float64,5,5), randn(Float64,5,5))
+
 @test gradtest(x -> sum(x, dims = (2, 3)), (3,4,5))
 @test gradtest(x -> sum(x, dims = 1), randn(Float64,2,3))
 @test gradtest(x -> sum(x, dims = [1,2]), randn(Float64,2,3))
 @test gradtest(x -> sum(x), randn(Float64,2,3))
+
 @test gradtest(x -> prod(x, dims=(2, 3)), (3,4,5))
+@test gradtest(x -> prod(x, dims=1), (3,4,5))
+@test gradtest(x -> prod(x, dims=1), (3,))
 @test gradtest(x -> prod(x), (3,4,5))
+@test gradtest(x -> prod(x), (3,))
+
+rz(dims...) = begin x = rand(dims...); x[2]=0; x end
+@test gradtest(x -> prod(x, dims=(2, 3)), rz(3,4,5))
+@test gradtest(x -> prod(x, dims=1), rz(3,4,5))
+@test gradtest(x -> prod(x, dims=1), rz(3,))
+@test gradtest(x -> prod(x), rz(3,4,5))
+@test gradtest(x -> prod(x), rz(3,))
+
+@test gradtest(x -> prod(x, dims=2), rz(3,4)') ## Adjoint tests fall-back methods
+@test gradtest(x -> prod(x, dims=2), rz(3,)')
+@test gradtest(x -> prod(x), rz(3,4)')
+@test gradtest(x -> prod(x), rz(3,)')
+
+@test gradtest(x -> cumprod(x, dims=2), (3,4,5))
+@test gradtest(x -> cumprod(x, dims=1), (3,)) 
+@test gradtest(x -> cumprod(x), (3,))
+
+@test gradtest(x -> cumprod(x, dims=2), rz(3,4,5))
+@test gradtest(x -> cumprod(x, dims=1), rz(3,)) 
+@test gradtest(x -> cumprod(x), rz(3,))
 
 @test gradtest(x -> softmax(x).*(1:3), 3)
 @test gradtest(x -> softmax(x).*(1:3), (3,5))