Merge pull request #328 from yuehhua/cartesianidx

Fix the output type of reverse_indices

src/utils.jl

@@ -34,10 +34,21 @@ function reverse_indices!(rev::AbstractArray, idx::AbstractArray)
     rev
 end
 
-function reverse_indices(idx::AbstractArray)
-    rev = Array{Vector{CartesianIndex}}(undef, maximum_dims(idx)...)
+"""
+    reverse_indices(idx)
+
+Return the reverse indices of `idx`. The indices of `idx` will be values, and values of `idx` will be index.
+
+# Arguments
+
+- `idx`: The indices to be reversed. Accepts array or cuarray of integer, tuple or `CartesianIndex`.
+"""
+function reverse_indices(idx::AbstractArray{<:Any,N}) where N
+    max_dims = maximum_dims(idx)
+    T = CartesianIndex{N}
+    rev = Array{Vector{T}}(undef, max_dims...)
     for i in eachindex(rev)
-        rev[i] = CartesianIndex[]
+        rev[i] = T[]
     end
     return reverse_indices!(rev, idx)
 end

test/utils.jl

@@ -24,13 +24,16 @@ end
            4 2 1 3;
            3 5 5 3]
     @test NNlib.reverse_indices(idx) == res
+    @test NNlib.reverse_indices(idx) isa typeof(res)
     idx = [(1,) (2,) (3,) (4,);
            (4,) (2,) (1,) (3,);
            (3,) (5,) (5,) (3,)]
     @test NNlib.reverse_indices(idx) == res
+    @test NNlib.reverse_indices(idx) isa typeof(res)
     idx = CartesianIndex.(
         [(1,) (2,) (3,) (4,);
         (4,) (2,) (1,) (3,);
         (3,) (5,) (5,) (3,)])
     @test NNlib.reverse_indices(idx) == res
+    @test NNlib.reverse_indices(idx) isa typeof(res)
 end