cleanup; getindex

FluxML · Mar 5, 2021 · 7307904 · 7307904
1 parent 920d7d1
commit 7307904
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -3,7 +3,6 @@ uuid = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 version = "0.7.14"
 
 [deps]
-BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

diff --git a/src/NNlib.jl b/src/NNlib.jl
@@ -7,7 +7,7 @@ import ChainRulesCore: rrule
 using Base.Broadcast: broadcasted
 using Statistics: mean
 
-const IntOrIntTuple = Union{Integer, NTuple{N,Integer} where N}
+const IntOrIntTuple = Union{Integer, NTuple{N,<:Integer} where N}
 const Numeric = Union{AbstractArray{<:T}, T} where {T<:Number}
 
 # Include APIs

diff --git a/src/gather.jl b/src/gather.jl
@@ -73,3 +73,12 @@ function gather(src::AbstractArray{Tsrc, Nsrc},
     dst = similar(src, Tsrc, dstsize)
     return gather!(dst, src, idx)
 end
+
+# Simple implementation with getindex for integer array.
+# Perf equivalent to the one above (which can also handle the integer case)
+# leave it here to show the simple connection with getindex.
+function gather(src::AbstractArray{Tsrc, Nsrc}, 
+                idx::AbstractArray{<:Integer}) where {Tsrc, Nsrc}
+    colons = ntuple(i -> Colon(), Nsrc-1)
+    return src[colons..., idx]
+end