Dimension mismatch in mapslices for TrackedArrays

[baggepinnen]

Ther following use of mapslices tries to call a non-existing Array constructor:

using Flux, CuArrays
julia> Z = gpu(Conv((3,3), 1=>1))(gpu(randn(5,5,1,1)))
Tracked 3×3×1×1 CuArray{Float32,4}:
[:, :, 1, 1] =
  1.85926   -4.03764  -1.79426
  0.348474  -3.98205  -1.1394 
 -0.766892  -1.29553  -1.96865

julia> Z = mapslices(Z, dims=(1,2,4)) do Zs
           Zs
       end
ERROR: MethodError: no method matching CuArray{Float32,4}(::UndefInitializer, ::Tuple{Int64,Int64,Int64})
Closest candidates are:
  CuArray{Float32,4}(::UndefInitializer, ::Tuple{Vararg{Int64,N}}) where {T, N} at /home/fredrikb/.julia/packages/CuArrays/wXQp8/src/array.jl:38
  CuArray{Float32,4}(::CUDAdrv.Mem.Buffer, ::Tuple{Vararg{Int64,N}}; offset, own) where {T, N} at /home/fredrikb/.julia/packages/CuArrays/wXQp8/src/array.jl:11
  CuArray{Float32,4}(::UndefInitializer, ::Integer...) where {T, N} at /home/fredrikb/.julia/packages/CuArrays/wXQp8/src/array.jl:42
  ...
Stacktrace:
 [1] similar(::Type{CuArray{Float32,4}}, ::Tuple{Int64,Int64,Int64}) at ./abstractarray.jl:672
 [2] convert(::Type{CuArray{Float32,4}}, ::Array{Float32,3}) at /home/fredrikb/.julia/packages/GPUArrays/J4c3Q/src/construction.jl:80
 [3] convert(::Type{CuArray{Float32,4}}, ::TrackedArray{…,CuArray{Float32,3}}) at /home/fredrikb/.julia/packages/GPUArrays/J4c3Q/src/construction.jl:71
 [4] gpu_convert at /home/fredrikb/.julia/packages/GPUArrays/J4c3Q/src/indexing.jl:119 [inlined]
 [5] _unsafe_setindex!(::IndexLinear, ::CuArray{Float32,4}, ::TrackedArray{…,CuArray{Float32,3}}, ::Base.OneTo{Int64}, ::Vararg{Union{Real, AbstractArray},N} where N) at /home/fredrikb/.julia/packages/GPUArrays/J4c3Q/src/indexing.jl:131
 [6] _setindex!(::IndexLinear, ::CuArray{Float32,4}, ::TrackedArray{…,CuArray{Float32,3}}, ::Base.OneTo{Int64}, ::Base.OneTo{Int64}, ::Vararg{Union{Real, AbstractArray},N} where N) at ./multidimensional.jl:712
 [7] setindex!(::CuArray{Float32,4}, ::TrackedArray{…,CuArray{Float32,3}}, ::Base.OneTo{Int64}, ::Base.OneTo{Int64}, ::Int64, ::Vararg{Any,N} where N) at ./abstractarray.jl:1074
 [8] concatenate_setindex!(::CuArray{Float32,4}, ::TrackedArray{…,CuArray{Float32,3}}, ::Base.OneTo{Int64}, ::Vararg{Any,N} where N) at ./abstractarray.jl:2060
 [9] #mapslices#109(::Tuple{Int64,Int64,Int64}, ::typeof(mapslices), ::getfield(Main, Symbol("##21#22")), ::TrackedArray{…,CuArray{Float32,4}}) at ./abstractarray.jl:2026
 [10] (::getfield(Base, Symbol("#kw##mapslices")))(::NamedTuple{(:dims,),Tuple{Tuple{Int64,Int64,Int64}}}, ::typeof(mapslices), ::Function, ::TrackedArray{…,CuArray{Float32,4}}) at ./none:0
 [11] top-level scope at none:0

whereas for a non-tracked array, it works

Z = CuArray{Float32,4}(undef, (5,5,1,1));
Z = mapslices(Z, dims=(1,2,4)) do Zs
                  Zs
 end

[AzamatB]

Bump. I'm hitting this.

[mcabbott]

Perhaps the same issue as https://github.com/FluxML/Flux.jl/issues/741. I don’t think mapslices works at all for Tracker or Zygote right now, although the prceise error varies:

julia> gradient(x -> sum(mapslices(identity, x, dims=1)), ones(2,2))
ERROR: MethodError: no method matching Float64(::Tracker.TrackedReal{Float64})