Views built within MLUtils.kfolds scalar indexing error on gpu

[tom-plaa]

Not sure if this is the correct repository to report this.
I'm trying to do kfolds cross validation using my GPU and I seem to be running into a problem caused by the way MLUtils builds the views. Hoping I can get some help to go around this.

The typeof(Flux.gpu(something_selected_by_MLUtils_kfolds)) is SubArray{Float32, 2, CuArray{Float32, 2, CUDA.DeviceMemory}, Tuple{Base.Slice{Base.OneTo{Int64}}, CuArray{Int64, 1, CUDA.DeviceMemory}}, false}

MWE:

using CUDA
using Flux
using MLUtils


function init_nn(inputsize::Integer)
    midsize = round(Int, inputsize / 2)
    block = Chain(
        Dense(inputsize => midsize),
        Dense(midsize => inputsize),
    )
    return block
end


inputs = randn(Float32, 1000, 3000)
labels = randn(Float32, 1000, 3000)
a = init_nn(1000)
b = init_nn(1000) |> Flux.gpu
kfolds_collection = MLUtils.kfolds((inputs, labels); k=2)

# Works
println("Test 1: ", Flux.mae(a(inputs), labels))
# Works
println("Test 2: ", Flux.mae(b(Flux.gpu(inputs)), Flux.gpu(labels)))

for (k, ((x_train, y_train), (x_val, y_val))) in enumerate(kfolds_collection)
    # Works
    println("Test fold $k, cpu train sets ", Flux.mae(a(x_train), y_train))
    # Works
    println("Test fold $k, cpu val sets ", Flux.mae(a(x_val), y_val))

    # Errors
    println("Test fold $k, gpu train sets ", Flux.mae(b(Flux.gpu(x_train)), Flux.gpu(y_train)))
    # Errors
    println("Test fold $k, gpu val sets ", Flux.mae(b(Flux.gpu(x_val)), Flux.gpu(y_val)))
end

The output:

Test 1: 1.0954581
Test 2: 1.0958811
Test fold 1, cpu train sets 1.0954655
Test fold 1, cpu val sets 1.0954509
ERROR: LoadError: Scalar indexing is disallowed.
Invocation of getindex resulted in scalar indexing of a GPU array.
This is typically caused by calling an iterating implementation of a method.
Such implementations *do not* execute on the GPU, but very slowly on the CPU,
and therefore should be avoided.

If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question.
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] errorscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:151
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:124
  [4] assertscalar(op::String)
    @ GPUArraysCore ~/.julia/packages/GPUArraysCore/aNaXo/src/GPUArraysCore.jl:112
  [5] getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:50 [inlined]
  [6] scalar_getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:36 [inlined]
  [7] _getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:19 [inlined]
  [8] getindex
    @ ~/.julia/packages/GPUArrays/u6tui/src/host/indexing.jl:17 [inlined]
  [9] _generic_matmatmul!(C::CuArray{…}, tA::Char, tB::Char, A::CuArray{…}, B::SubArray{…}, _add::LinearAlgebra.MulAddMul{…})
    @ LinearAlgebra ~/.julia/juliaup/julia-1.10.10+0.x64.linux.gnu/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:816
 [10] generic_matmatmul!(C::CuArray{…}, tA::Char, tB::Char, A::CuArray{…}, B::SubArray{…}, _add::LinearAlgebra.MulAddMul{…})
    @ LinearAlgebra ~/.julia/juliaup/julia-1.10.10+0.x64.linux.gnu/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:785
 [11] mul!
    @ ~/.julia/juliaup/julia-1.10.10+0.x64.linux.gnu/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:263 [inlined]
 [12] mul!
    @ ~/.julia/juliaup/julia-1.10.10+0.x64.linux.gnu/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:237 [inlined]
 [13] *
    @ ~/.julia/juliaup/julia-1.10.10+0.x64.linux.gnu/share/julia/stdlib/v1.10/LinearAlgebra/src/matmul.jl:106 [inlined]
 [14] (::Dense{typeof(identity), CuArray{…}, CuArray{…}})(x::SubArray{Float32, 2, CuArray{…}, Tuple{…}, false})
    @ Flux ~/.julia/packages/Flux/uRn8o/src/layers/basic.jl:199
 [15] macro expansion
    @ ~/.julia/packages/Flux/uRn8o/src/layers/basic.jl:68 [inlined]
 [16] _applychain
    @ ~/.julia/packages/Flux/uRn8o/src/layers/basic.jl:68 [inlined]
 [17] (::Chain{Tuple{…}})(x::SubArray{Float32, 2, CuArray{…}, Tuple{…}, false})
    @ Flux ~/.julia/packages/Flux/uRn8o/src/layers/basic.jl:65

Versions:

Version 1.10.10 (2025-06-27)

[052768ef] CUDA v5.8.2
[587475ba] Flux v0.16.5
[f1d291b0] MLUtils v0.4.8
[02a925ec] cuDNN v1.4.3

[tom-plaa]

By the way, I just noticed that if I do something like Flux.gpu(collect(kfolds_subarray)) then it works, but I have no idea if I am sacrificing performance.