Fix sum of tracked reals in Turing (#173)

* fix sum for tracked reals

* fix sum implementation

* fix the fix

* pull value

* apply rule to trackedarrays only

Author: mohdibntarek

src/derivatives/linalg/reductions.jl

@@ -5,24 +5,30 @@
 # basic sum #
 #-----------#
 
-function Base.sum(x::TrackedArray{V,D}) where {V,D}
+function Base.sum(x::TrackedArray{V,D}; dims=:) where {V,D}
     tp = tape(x)
-    out = track(sum(value(x)), D, tp)
-    record!(tp, SpecialInstruction, sum, x, out)
+    out = track(sum(value(x), dims = dims), D, tp)
+    record!(tp, SpecialInstruction, sum, (x, dims), out)
     return out
 end
 
 @noinline function special_reverse_exec!(instruction::SpecialInstruction{typeof(sum)})
-    input = instruction.input
+    input, dims = instruction.input
     output = instruction.output
-    istracked(input) && increment_deriv!(input, deriv(output))
+    if istracked(input)
+        if dims === Colon()
+            increment_deriv!(input, deriv(output))
+        else
+            increment_deriv!(input, zero(value(input)) .+ deriv(output))
+        end
+    end
     unseed!(output)
     return nothing
 end
 
 @noinline function special_forward_exec!(instruction::SpecialInstruction{typeof(sum)})
-    input = instruction.input
-    value!(instruction.output, sum(value(input)))
+    input, dims = instruction.input
+    value!(instruction.output, sum(value(input); dims = dims))
     return nothing
 end
 

src/tracked.jl

@@ -222,7 +222,7 @@ unseed!(x::AbstractArray, i) = unseed!(x[i])
 # `forward_pass!`/`reverse_pass!`.
 capture(t::TrackedReal) = ifelse(hastape(t), t, value(t))
 capture(t::TrackedArray) = t
-capture(t::AbstractArray) = istracked(t) ?  map!(capture, similar(t), t) : copy(t)
+capture(t::AbstractArray) = istracked(t) ? map(capture, t) : copy(t)
 
 ########################
 # Conversion/Promotion #
@@ -415,6 +415,7 @@ Base.float(t::TrackedReal{V}) where {V<:AbstractFloat} = t
 
 Base.one(::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(one(V))
 Base.zero(::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(zero(V))
+Base.zero(::Type{<:TrackedReal{V,D}}) where {V,D} = TrackedReal{V,D,Nothing}(zero(V))
 
 Base.rand(::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(rand(V))
 Base.rand(rng::Random.AbstractRNG, ::Type{TrackedReal{V,D,O}}) where {V,D,O} = TrackedReal{V,D,O}(rand(rng, V))