Documentation Example Causes Scalar Indexing Error wit Metal

[lenianiva]

I'm running https://juliagraphs.org/GraphNeuralNetworks.jl/docs/GraphNeuralNetworks.jl/stable/ on Julia v1.11.6 and GNN v1.0.0 installed via add GraphNeuralNetworks

using GraphNeuralNetworks, Flux, Statistics, MLUtils
using Flux: DataLoader
import Metal

all_graphs = GNNGraph[]

for _ in 1:1000
    g = rand_graph(10, 40,
            ndata=(; x = randn(Float32, 16,10)),  # Input node features
            gdata=(; y = randn(Float32)))         # Regression target
    push!(all_graphs, g)
end

device = gpu

model = GNNChain(
    GCNConv(16 => 64),
    BatchNorm(64),     # Apply batch normalization on node features (nodes dimension is batch dimension)
    x -> relu.(x),
    GCNConv(64 => 64, relu),
    GlobalPool(mean),  # Aggregate node-wise features into graph-wise features
    Dense(64, 1)) |> device

opt = Flux.setup(Adam(1f-4), model)

train_graphs, test_graphs = MLUtils.splitobs(all_graphs, at=0.8)

train_loader = DataLoader(
    train_graphs,
    batchsize=32,
    shuffle=true, collate=true)
test_loader = DataLoader(
    test_graphs,
    batchsize=32,
    shuffle=false, collate=true)

loss(model, g::GNNGraph) = mean((vec(model(g, g.x)), g.y - g.y).^2)

loss(model, loader) = mean(loss(model, g |> device) for g in loader)

for epoch in 1:1
    for g in train_loader
        g = g |> device
        grad = gradient(model -> loss(model, g), model)
        Flux.update!(opt, model, grad[1])
    end

    @info (; epoch, train_loss=loss(model, train_loader), test_loss=loss(model, test_loader))
end

When I ran this code, I got

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/ZRk7Q/src/host/indexing.jl:50 [inlined]
  [6] sparse!(I::Metal.MtlVector{Int64, Metal.PrivateStorage}, J::Metal.MtlVector{Int64, Metal.PrivateStorage}, V::Metal.MtlVector{Int64, Metal.PrivateStorage}, m::Int64, n::Int64, combine::typeof(+), klasttouch::Vector{Int64}, csrrowptr::Vector{Int64}, csrcolval::Vector{Int64}, csrnzval::Vector{Int64}, csccolptr::Vector{Int64}, cscrowval::Vector{Int64}, cscnzval::Vector{Int64})
    @ SparseArrays ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:1172
  [7] sparse(I::Metal.MtlVector{Int64, Metal.PrivateStorage}, J::Metal.MtlVector{Int64, Metal.PrivateStorage}, V::Metal.MtlVector{Int64, Metal.PrivateStorage}, m::Int64, n::Int64, combine::Function)
    @ SparseArrays ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:1093
  [8] sparse
    @ ~/.julia/juliaup/julia-1.11.6+0.aarch64.apple.darwin14/share/julia/stdlib/v1.11/SparseArrays/src/sparsematrix.jl:1328 [inlined]
  [9] #to_sparse#333
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/convert.jl:231 [inlined]
 [10] to_sparse
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/convert.jl:221 [inlined]
 [11] #adjacency_matrix#82
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/query.jl:227 [inlined]
 [12] adjacency_matrix
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/query.jl:220 [inlined]
 [13] #rrule#87
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/query.jl:264 [inlined]
 [14] rrule
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/query.jl:262 [inlined]
 [15] rrule
    @ ~/.julia/packages/ChainRulesCore/Vsbj9/src/rules.jl:144 [inlined]
 [16] chain_rrule_kw
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/chainrules.jl:246 [inlined]
 [17] macro expansion
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0 [inlined]
 [18] _pullback(::Zygote.Context{false}, ::typeof(Core.kwcall), ::@NamedTuple{weighted::Bool}, ::typeof(adjacency_matrix), ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::Type{Int64})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:81
 [19] adjacency_matrix
    @ ~/.julia/packages/GNNGraphs/Ldvz4/src/query.jl:220 [inlined]
 [20] _pullback(::Zygote.Context{false}, ::typeof(Core.kwcall), ::@NamedTuple{weighted::Bool}, ::typeof(adjacency_matrix), ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [21] propagate
    @ ~/.julia/packages/GNNlib/StdFP/src/msgpass.jl:216 [inlined]
 [22] _pullback(::Zygote.Context{false}, ::typeof(propagate), ::typeof(copy_xj), ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::typeof(+), ::Nothing, ::Metal.MtlMatrix{Float32, Metal.PrivateStorage}, ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [23] #propagate#1
    @ ~/.julia/packages/GNNlib/StdFP/src/msgpass.jl:72 [inlined]
 [24] _pullback(::Zygote.Context{false}, ::GNNlib.var"##propagate#1", ::Nothing, ::Metal.MtlMatrix{Float32, Metal.PrivateStorage}, ::Nothing, ::typeof(propagate), ::typeof(copy_xj), ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::typeof(+))
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [25] propagate
    @ ~/.julia/packages/GNNlib/StdFP/src/msgpass.jl:71 [inlined]
 [26] _pullback(::Zygote.Context{false}, ::typeof(Core.kwcall), ::@NamedTuple{xj::Metal.MtlMatrix{Float32, Metal.PrivateStorage}}, ::typeof(propagate), ::typeof(copy_xj), ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::typeof(+))
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [27] gcn_conv
    @ ~/.julia/packages/GNNlib/StdFP/src/layers/conv.jl:65 [inlined]
 [28] _pullback(::Zygote.Context{false}, ::typeof(gcn_conv), ::GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::Metal.MtlMatrix{Float32, Metal.PrivateStorage}, ::Nothing, ::GraphNeuralNetworks.var"#21#23", ::Nothing)
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [29] #_#20
    @ ~/.julia/packages/GraphNeuralNetworks/nauyk/src/layers/conv.jl:103 [inlined]
 [30] GCNConv (repeats 2 times)
    @ ~/.julia/packages/GraphNeuralNetworks/nauyk/src/layers/conv.jl:99 [inlined]
 [31] _applylayer
    @ ~/.julia/packages/GraphNeuralNetworks/nauyk/src/layers/basic.jl:152 [inlined]
 [32] _pullback(::Zygote.Context{false}, ::typeof(GraphNeuralNetworks._applylayer), ::GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::Metal.MtlMatrix{Float32, Metal.PrivateStorage})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [33] _applychain
    @ ~/.julia/packages/GraphNeuralNetworks/nauyk/src/layers/basic.jl:138 [inlined]
 [34] _pullback(::Zygote.Context{false}, ::typeof(GraphNeuralNetworks._applychain), ::Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}, ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::Metal.MtlMatrix{Float32, Metal.PrivateStorage})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [35] GNNChain
    @ ~/.julia/packages/GraphNeuralNetworks/nauyk/src/layers/basic.jl:124 [inlined]
 [36] _pullback(::Zygote.Context{false}, ::GNNChain{Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}}, ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}}, ::Metal.MtlMatrix{Float32, Metal.PrivateStorage})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [37] loss
    @ ~/Projects/simulations/GNNExp/src/minimal-gnn.jl:37 [inlined]
 [38] _pullback(::Zygote.Context{false}, ::typeof(loss), ::GNNChain{Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}}, ::GNNGraph{Tuple{Metal.MtlVector{Int64, Metal.PrivateStorage}, Metal.MtlVector{Int64, Metal.PrivateStorage}, Nothing}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [39] #5
    @ ~/Projects/simulations/GNNExp/src/minimal-gnn.jl:44 [inlined]
 [40] _pullback(ctx::Zygote.Context{false}, f::var"#5#6", args::GNNChain{Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface2.jl:0
 [41] pullback(f::Function, cx::Zygote.Context{false}, args::GNNChain{Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:96
 [42] pullback
    @ ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:94 [inlined]
 [43] gradient(f::Function, args::GNNChain{Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}})
    @ Zygote ~/.julia/packages/Zygote/55SqB/src/compiler/interface.jl:153
 [44] #gradient#1
    @ ~/.julia/packages/Flux/uRn8o/src/gradient.jl:44 [inlined]
 [45] gradient(f::Function, args::GNNChain{Tuple{GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(identity)}, BatchNorm{typeof(identity), Metal.MtlVector{Float32, Metal.PrivateStorage}, Float32, Metal.MtlVector{Float32, Metal.PrivateStorage}}, var"#1#2", GCNConv{Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}, typeof(relu)}, GlobalPool{typeof(mean)}, Dense{typeof(identity), Metal.MtlMatrix{Float32, Metal.PrivateStorage}, Metal.MtlVector{Float32, Metal.PrivateStorage}}}})
    @ Flux ~/.julia/packages/Flux/uRn8o/src/gradient.jl:31

The error seems to come from a layer in the GNN. How can I solve this issue? Do I need to enable scalar indexing?

[CarloLucibello]

Unfortunately Metal is not fully supported yet. With respect to this example, #619 improves things but the example still fails due to GlobalPool using scatter, which is not supported on Metal.