Add scatter operations

src/NNlib.jl

@@ -4,7 +4,9 @@ using Pkg
 using Requires
 using ChainRulesCore
 using Base.Broadcast: broadcasted
+using Statistics: mean
 
+const IntOrTuple = Union{Integer,Tuple}
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
 # Include APIs
@@ -31,6 +33,8 @@ include("conv.jl")
 include("conv_bias_act.jl")
 include("pooling.jl")
 include("upsample.jl")
+include("utils.jl")
+include("scatter.jl")
 
 ## Include implementations
 include("impl/padding_edges.jl")

src/scatter.jl

@@ -0,0 +1,189 @@
+export scatter!, scatter
+
+## Scatter API
+#   - Scatter:
+#     - scatter(op, src, idx)
+#     - scatter!(op, dst, src, idx)
+#   - Scatter destination backpropagation
+#     - ∇scatter_dst!
+#   - Scatter source backpropagation
+#     - ∇scatter_src
+#     - ∇scatter_src!
+#
+
+function _check_dims(Ndst, Nsrc, N, Nidx)
+    @assert Ndst - N == Nsrc - Nidx
+    dims = Ndst - N
+    if dims < 0
+        throw(ArgumentError("dims must be non-negative but got dims=$dims"))
+    end
+    return dims
+end
+
+_check_input(idx::AbstractArray{<:Integer}, arr) = checkbounds(arr, minimum(idx):maximum(idx))
+
+function _check_input(idx::AbstractArray{<:Tuple}, arr)
+    pairs = map(xs -> Base.OneTo(maximum(xs)), zip(idx...))
+    checkbounds(arr, pairs...)
+end
+
+function _check_output(idx::AbstractArray{<:IntOrTuple}, arr, src, dims)
+    pre_dims = axes(src)[1:dims]
+    post_dims = Base.OneTo.(maximum_dims(idx))
+    checkbounds(arr, pre_dims..., post_dims...)
+end
+
+
+"""
+    scatter!(op, dst, src, idx)
+
+Scatter operation, which scatters data in `src` and assigns to `dst` according to `idx`.
+With the data going to the same place, specified operation is applied on to reduce data.
+For each index `k` in `idx`, accumulate values in `dst` according to
+
+    dst[idx[k]...] = (op).(dst[idx[k]...], src[k...])
+
+# Arguments
+- `op`: operations to be applied on `dst` and `src`, e.g. `+`, `-`, `*`, `/`, `max` and `min`.
+- `dst`: the destination for `src` to aggregate to. This argument will be mutated.
+- `src`: the source data for aggregating.
+- `idx`: the mapping for aggregation from source (index) to destination (value).
+The index of `idx` is corresponding to the index of `src` and the value of `idx` is
+corresponding to the index of `dst`. The value of `idx` can be `Int` or `Tuple` type.
+
+    dst[:, ..., idx[k]...] = (op).(dst[:, ..., idx[k]...], src[:, ..., k...])
+"""
+function scatter!(op,
+                  dst::AbstractArray{Tdst,Ndst},
+                  src::AbstractArray{Tsrc,Nsrc},
+                  idx::AbstractArray{<:Integer,Nidx}) where {Tdst<:Real,Tsrc<:Real,Ndst,Nsrc,Nidx}
+    dims = _check_dims(Ndst, Nsrc, 1, Nidx)
+    @boundscheck _check_output(idx, dst, src, dims)
+    @boundscheck _check_input(idx, src)
+    T = promote_type(Tdst, Tsrc)
+    scatter!(op, T.(dst), src, idx, Val(dims))
+end
+
+function scatter!(op,
+                  dst::AbstractArray{Tdst,Ndst},
+                  src::AbstractArray{Tsrc,Nsrc},
+                  idx::AbstractArray{NTuple{N,Int},Nidx}) where {Tdst<:Real,Tsrc<:Real,Ndst,Nsrc,N,Nidx}
+    dims = _check_dims(Ndst, Nsrc, N, Nidx)
+    @boundscheck _check_output(idx, dst, src, dims)
+    @boundscheck _check_input(idx, src)
+    T = promote_type(Tdst, Tsrc)
+    scatter!(op, T.(dst), src, idx, Val(dims))
+end
+
+function scatter!(op, dst::AbstractArray{Tdst}, src::AbstractArray{Tsrc}, idx::AbstractArray{<:IntOrTuple},
+                  dims::Val{N}) where {Tdst<:Real,Tsrc<:Real,N}
+    colons = Base.ntuple(_->Colon(), N)
+    @simd for k in CartesianIndices(idx)
+        dst_v = view(dst, colons..., idx[k]...)
+        src_v = view(src, colons..., k)
+        @inbounds dst_v .= (op).(dst_v, src_v)
+    end
+    dst
+end
+
+"""
+    scatter!(mean, dst, src, idx)
+
+Scatter mean operation, which scatters data in `src` and assigns to `dst` according to `idx`.
+With the data going to the same place, mean is applied on to reduce data.
+For each index `k` in `idx`, accumulate values in `dst` according to
+
+    dst[idx[k]...] = dst[idx[k]...] + mean.(src[k...])
+
+# Arguments
+- `dst`: the destination for `src` to aggregate to. This argument will be mutated.
+- `src`: the source data for aggregating.
+- `idx`: the mapping for aggregation from source (index) to destination (value).
+The index of `idx` is corresponding to the index of `src` and the value of `idx` is
+corresponding to the index of `dst`. The value of `idx` can be `Int` or `Tuple` type.
+
+    dst[:, ..., idx[k]...] = (op).(dst[:, ..., idx[k]...], src[:, ..., k...])
+"""
+function scatter!(op::typeof(mean),
+                  dst::AbstractArray{Tdst,Ndst},
+                  src::AbstractArray{Tsrc,Nsrc},
+                  idx::AbstractArray{<:Integer,Nidx}) where {Tdst<:Real,Tsrc<:Real,Ndst,Nsrc,Nidx}
+    Ns = scatter!(+, zero(dst), one.(src), idx)
+    dst_ = scatter!(+, zero(dst), src, idx)
+    dst .+= safe_div.(dst_, Ns)
+    return dst
+end
+
+function scatter!(op::typeof(mean),
+                  dst::AbstractArray{Tdst,Ndst},
+                  src::AbstractArray{Tsrc,Nsrc},
+                  idx::AbstractArray{NTuple{N,Int},Nidx}) where {Tdst<:Real,Tsrc<:Real,Ndst,Nsrc,N,Nidx}
+    Ns = scatter!(+, zero(dst), one.(src), idx)
+    dst_ = scatter!(+, zero(dst), src, idx)
+    dst .+= safe_div.(dst_, Ns)
+    return dst
+end
+
+
+"""
+    scatter(op, src, idx)
+
+Scatter operation, which applies specified operation on `src` according to `idx`
+and gives an new array `dst`.
+For each index `k` in `idx`, accumulate values in `dst` according to
+
+    dst[idx[k]...] = (op).(src[k...])
+
+# Arguments
+- `op`: operations to be applied on `dst` and `src`, e.g. `+`, `-`, `*`, `/`, `max` and `min`.
+- `src`: the source data for aggregating.
+- `idx`: the mapping for aggregation from source (index) to destination (value).
+The index of `idx` is corresponding to the index of `src` and the value of `idx` is
+corresponding to the index of `dst`. The value of `idx` can be `Int` or `Tuple` type.
+"""
+function scatter end
+
+for op in [+, -]
+    @eval function scatter(op::typeof($op),
+                           src::AbstractArray{T,Nsrc},
+                           idx::AbstractArray{<:IntOrTuple,Nidx}) where {T<:Real,Nsrc,Nidx}
+        dims = Nsrc - Nidx
+        dst = zeros(T, size(src)[1:dims]..., maximum_dims(idx)...)
+        scatter!(op, dst, src, idx)
+    end
+end
+
+for op in [*, /]
+    @eval function scatter(op::typeof($op),
+                           src::AbstractArray{T,Nsrc},
+                           idx::AbstractArray{<:IntOrTuple,Nidx}) where {T<:Real,Nsrc,Nidx}
+        dims = Nsrc - Nidx
+        dst = ones(T, size(src)[1:dims]..., maximum_dims(idx)...)
+        scatter!(op, dst, src, idx)
+    end
+end
+
+function scatter(op::typeof(max),
+                 src::AbstractArray{T,Nsrc},
+                 idx::AbstractArray{<:IntOrTuple,Nidx}) where {T<:Real,Nsrc,Nidx}
+    dims = Nsrc - Nidx
+    dst = fill(typemin(T), size(src)[1:dims]..., maximum_dims(idx)...)
+    scatter!(op, dst, src, idx)
+end
+
+function scatter(op::typeof(min),
+                 src::AbstractArray{T,Nsrc},
+                 idx::AbstractArray{<:IntOrTuple,Nidx}) where {T<:Real,Nsrc,Nidx}
+    dims = Nsrc - Nidx
+    dst = fill(typemax(T), size(src)[1:dims]..., maximum_dims(idx)...)
+    scatter!(op, dst, src, idx)
+end
+
+function scatter(op::typeof(mean),
+                 src::AbstractArray{T,Nsrc},
+                 idx::AbstractArray{<:IntOrTuple,Nidx}) where {T<:Real,Nsrc,Nidx}
+    FT = float(T)
+    dims = Nsrc - Nidx
+    dst = zeros(FT, size(src)[1:dims]..., maximum_dims(idx)...)
+    scatter!(op, dst, FT.(src), idx)
+end

src/utils.jl

@@ -0,0 +1,42 @@
+"""
+    safe_div(x, y)
+
+Safely divide `x` by `y`. If `y` is zero, return `x` directly.
+"""
+safe_div(x, y) = ifelse(iszero(y), x, x/y)
+
+"""
+    maximum_dims(dims)
+
+Return the maximum value for each dimension. An array of dimensions `dims` is accepted.
+The maximum of each dimension in the element is computed.
+"""
+maximum_dims(dims::AbstractArray{<:Integer}) = (maximum(dims), )
+
+function maximum_dims(dims::AbstractArray{<:Tuple})
+    Tuple(maximum(xs) for xs in zip(dims...))
+end
+
+function reverse_indices(X::Array{T}) where T
+    Y = Dict{T,Vector{CartesianIndex}}()
+    @inbounds for (ind, val) = pairs(X)
+        Y[val] = get(Y, val, CartesianIndex[])
+        push!(Y[val], ind)
+    end
+    Y
+end
+
+function count_indices(idx::AbstractArray, N)
+    counts = zero.(idx)
+    @inbounds for i = 1:N
+        counts += sum(idx.==i) * (idx.==i)
+    end
+    counts
+end
+
+function divide_by_counts!(xs, idx::AbstractArray, N)
+    counts = count_indices(idx, N)
+    @inbounds for ind = CartesianIndices(counts)
+        view(xs, :, ind) ./= counts[ind]
+    end
+end

test/runtests.jl

@@ -41,3 +41,11 @@ end
 @testset "Upsampling" begin
     include("upsample.jl")
 end
+
+@testset "Scatter" begin
+    include("scatter.jl")
+end
+
+@testset "Utilities" begin
+    include("utils.jl")
+end