Zero-sized Array after indexing causes exception during gradient

[lenianiva]

In some cases, when there are arrays a, i, with neither empty, but a[i] is empty, calculating the gradient gives errors:

using NNlib, Flux, Random, Statistics, Zygote

major = 7

to_ordered_index(i::AbstractArray) = Zygote.ignore(() -> CartesianIndex.(1:length(i), i))
function loss_of(model, a::AbstractArray, i::AbstractArray)
    mask = i .!= 0
    #iszero(mask) && return 0
    i_ = i[mask]
    Zygote.ignore() do
        @info "i: $i, i_: $i_"
    end
    j = to_ordered_index(i_)
    b = model(a)
    c = b[:,mask]'[j]
    return mean(c .^ 2)
end

model = Dense(major => major) |> f32 |> gpu
opt = Flux.setup(Adam(1e-3), model)
a = ones((major, 5)) |> f32 |> gpu
i = [0, 0] |> gpu

gs = Flux.gradient(model) do m
    loss_of(m, a, i)
end

Flux.update!(opt, a, gs)

the error is [edit: extended]

julia> gs = Flux.gradient(model) do m
           loss_of(m, a, i)
       end
[ Info: i: [0, 0], i_: Int64[]
ERROR: BoundsError: attempt to access 7×5 Matrix{Float32} at index [1:7, 2-element BitVector]
Stacktrace:
  [1] throw_boundserror(A::Matrix{Float32}, I::Tuple{Base.Slice{Base.OneTo{Int64}}, Base.LogicalIndex{Int64, BitVector}})
    @ Base ./essentials.jl:14
  [2] checkbounds
    @ ./abstractarray.jl:699 [inlined]
  [3] _getindex
    @ ./multidimensional.jl:914 [inlined]
  [4] getindex
    @ ./abstractarray.jl:1312 [inlined]
  [5] rrule
    @ ~/.julia/packages/ChainRules/14CDN/src/rulesets/Base/indexing.jl:63 [inlined]
  [6] rrule
    @ ~/.julia/packages/ChainRulesCore/Vsbj9/src/rules.jl:138 [inlined]
  [7] chain_rrule
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/chainrules.jl:234 [inlined]
  [8] macro expansion
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0 [inlined]
  [9] _pullback
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:81 [inlined]
 [10] loss_of
    @ ./REPL[35]:10 [inlined]
 [11] _pullback(::Zygote.Context{…}, ::typeof(loss_of), ::Dense{…}, ::Matrix{…}, ::Vector{…})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [12] #13
    @ ./REPL[40]:2 [inlined]
 [13] _pullback(ctx::Zygote.Context{false}, f::var"#13#14", args::Dense{typeof(identity), Matrix{Float32}, Vector{Float32}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [14] pullback(f::Function, cx::Zygote.Context{false}, args::Dense{typeof(identity), Matrix{Float32}, Vector{Float32}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:96
 [15] pullback
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:94 [inlined]
 [16] gradient(f::Function, args::Dense{typeof(identity), Matrix{Float32}, Vector{Float32}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:153
 [17] #gradient#1
    @ ~/.julia/packages/Flux/WMUyh/src/gradient.jl:44 [inlined]
 [18] gradient(f::Function, args::Dense{typeof(identity), Matrix{Float32}, Vector{Float32}})
    @ Flux ~/.julia/packages/Flux/WMUyh/src/gradient.jl:31
 [19] top-level scope
    @ REPL[40]:1
Some type information was truncated. Use `show(err)` to see complete types.

Uncomment the #iszero(mask) line for a temporary solution. Although the code says |> gpu, the error occurs on CPU as well.

Note that the behaviour of this code is somewhat unstable. It seems to be accessing already free'd memory, so the same code could sometimes not end in an error on a long-running Julia session.

[lenianiva]

Another example where the out of bounds error only happens on GPU. This array is not zero sized

using CUDA, cuDNN
using NNlib, Flux, Random, Statistics

dim_entries = 5
dim_embedding = 7

function f(
    a::AbstractArray,
    i::AbstractArray,
    target::AbstractArray,
    model,
)
    iszero(i) && return zeros(eltype(a)) |> gpu

    mask = i .!= 0
    lhs = a[:, mask]
    rhs = model(i[mask])
    z = dropdims(sum(lhs .* rhs, dims = 1), dims = 1)
    return z - target[mask]
end

model = Embedding(dim_entries => dim_embedding) |> f32 |> gpu
opt = Flux.setup(Adam(1e-3), model)

i = [3,4,5,0,2] |> gpu
a = rand(dim_embedding, length(i)) |> gpu
target = rand(length(i)) |> gpu

gs = Flux.gradient(model) do m
    sum(f(a, i, target, model))
end

Flux.update!(opt, a, gs)

gives

ERROR: CUDA error: an illegal memory access was encountered (code 700, ERROR_ILLEGAL_ADDRESS)
Stacktrace:
  [1] throw_api_error(res::CUDA.cudaError_enum)
    @ CUDA ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/lib/cudadrv/libcuda.jl:30
  [2] check
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/lib/cudadrv/libcuda.jl:37 [inlined]
  [3] cuMemAllocFromPoolAsync
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/GPUToolbox/src/ccalls.jl:33 [inlined]
  [4] alloc(::Type{CUDA.DeviceMemory}, bytesize::Int64; async::Bool, stream::CuStream, pool::CuMemoryPool)
    @ CUDA ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/lib/cudadrv/memory.jl:71
  [5] alloc
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/lib/cudadrv/memory.jl:61 [inlined]
  [6] (::CUDA.var"#_pool_alloc##0#_pool_alloc##1"{Int64, CuMemoryPool, @NamedTuple{device::CuDevice, context::CuContext, stream::CuStream, math_mode::CUDA.MathMode, math_precision::Symbol}})()
    @ CUDA ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/memory.jl:645
  [7] retry_reclaim
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/memory.jl:434 [inlined]
  [8] _pool_alloc
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/memory.jl:640 [inlined]
  [9] macro expansion
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/memory.jl:623 [inlined]
 [10] macro expansion
    @ ./timing.jl:461 [inlined]
 [11] pool_alloc
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/memory.jl:622 [inlined]
 [12] (::CUDA.var"#650#651"{CUDA.DeviceMemory, Int64})()
    @ CUDA ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/array.jl:92
 [13] cached_alloc(f::CUDA.var"#650#651"{CUDA.DeviceMemory, Int64}, key::Tuple{UnionAll, CuDevice, DataType, Int64})
    @ GPUArrays ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/GPUArrays/src/host/alloc_cache.jl:36
 [14] CuArray{Float32, 2, CUDA.DeviceMemory}(::UndefInitializer, dims::Tuple{Int64, Int64})
    @ CUDA ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/array.jl:91
 [15] similar
    @ ./abstractarray.jl:867 [inlined]
 [16] similar
    @ ./abstractarray.jl:866 [inlined]
 [17] similar
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/CUDA/src/broadcast.jl:25 [inlined]
 [18] similar
    @ ./broadcast.jl:227 [inlined]
 [19] copy
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/GPUArrays/src/host/broadcast.jl:29 [inlined]
 [20] materialize
    @ ./broadcast.jl:894 [inlined]
 [21] adjoint
    @ ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/Zygote/src/lib/broadcast.jl:93 [inlined]
 [22] _pullback(__context__::Zygote.Context{false}, 637::typeof(Base.Broadcast.broadcasted), 638::typeof(*), x::CuArray{Float32, 2, CUDA.DeviceMemory}, y::CuArray{Float32, 2, CUDA.DeviceMemory})
    @ Zygote ~/nix/store/jpiwcw5d7ylavwbf2qqyxlhzw5dp24yx-trexp-cuda-load-path/ZygoteRules/src/adjoint.jl:67
 [23] f
    @ ./REPL[5]:7 [inlined]

[mcabbott]

Any chance you can boil these down? There's probably one offending operation, your line [22] hints that it may just be x .* y with one array empty.

[lenianiva]

Any chance you can boil these down? There's probably one offending operation, your line [22] hints that it may just be x .* y with one array empty.

do you mean the first example with Dense or the second one with Embedding?

[mcabbott]

Worth shortening both, unless you think they are eventually the same error.

[22] referred to the second; I've now edited the first to include the full stacktrace.