Merge pull request #350 from FluxML/cl/scatter

simplify scatter mean
FluxML · Sep 7, 2021 · 20c110b · 20c110b · CarloLucibello · Sep 7, 2021
2 parents 5256d0b + 8962dc7
commit 20c110b
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 22 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "NNlib"
 uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
-version = "0.7.28"
+version = "0.7.29"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/gather.jl b/src/gather.jl
@@ -63,7 +63,6 @@ function gather(src::AbstractArray{Tsrc, Nsrc},
     return gather!(dst, src, idx)
 end
 
-
 ∇gather_src(Δ, src_size, idx) = scatter!(+, fill!(similar(Δ, eltype(Δ), src_size), 0), Δ, idx)
 
 function rrule(::typeof(gather!), dst::AbstractArray, src::AbstractArray, idx::AbstractArray)

diff --git a/src/scatter.jl b/src/scatter.jl
@@ -148,8 +148,17 @@ end
 function ∇scatter_src(::typeof(mean), Δ, dst,
                     src::AbstractArray{Tsrc,Nsrc}, 
                     idx::AbstractArray{Tidx,Nidx}) where {Tsrc,Tidx,Nsrc,Nidx}
-    dims = Nsrc - Nidx
-    divide_by_counts!(gather(Δ, idx), idx, dims)
+
+    M = typelength(Tidx)
+    num = gather(Δ, idx)
+    counts = fill!(similar(Δ, Int, size(Δ)[end-M+1:end]), 0)
+    scatter!(+, counts, fill!(similar(idx, Int), 1), idx)
+    den = gather(counts, idx)
+    # make num and den broadcast compatible
+    for i in 1:ndims(num)-ndims(den)
+        den = unsqueeze(den)
+    end
+    return safe_div.(num, den)
 end
 
 ∇scatter_src(op, Δ, dst, src, idx) =

diff --git a/src/utils.jl b/src/utils.jl
@@ -53,19 +53,4 @@ function reverse_indices(idx::AbstractArray{<:Any,N}) where N
     return reverse_indices!(rev, idx)
 end
 
-function count_indices(idx::AbstractArray)
-    counts = zero.(idx)
-    for i in unique(idx)
-        counts += sum(idx .== i) * (idx .== i)
-    end
-    return counts
-end
-
-function divide_by_counts!(xs, idx::AbstractArray, dims)
-    colons = Base.ntuple(_->Colon(), dims)
-    counts = count_indices(idx)
-    for i in CartesianIndices(counts)
-        view(xs, colons..., i) ./= counts[i]
-    end
-    return xs
-end
+unsqueeze(x) = reshape(x, 1, size(x)...) 
diff --git a/test/scatter.jl b/test/scatter.jl
@@ -68,8 +68,7 @@ res = Dict(
                          4. 4. 6. 5. 5.],
 )
 
-types = [UInt8,  UInt32, UInt128,
-         Int16, Int64, BigInt,
+types = [UInt8,  UInt32, Int64,
          Float16, Float32, Float64, BigFloat, Rational]
 
 @testset "scatter" begin
-Original file line number
+Diff line change
@@ Expand Up / @@ -63,7 +63,6 @@ function gather(src::AbstractArray{Tsrc, Nsrc}, @@
         return gather!(dst, src, idx)
     end
     ∇gather_src(Δ, src_size, idx) = scatter!(+, fill!(similar(Δ, eltype(Δ), src_size), 0), Δ, idx)
     function rrule(::typeof(gather!), dst::AbstractArray, src::AbstractArray, idx::AbstractArray)
@@ Expand Down @@