Documentation Example Causes Scalar Indexing Error wit Metal

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?

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.

@CarloLucibello

@CarloLucibello do you think adding a few of the ML-oriented operations provided by the Metal Performance Shaders might fill the gaps in the GNN API? For instance, MPS has specialized scatter operations which are probably already wrapped by Metal.jl

I ask because I'd be curious to take a stab at it, but I'm not sure it would be enough to bring Metal closer to being usable

@alonsoC1s

@alonsoC1s adding missing operators in Metal.jl, NNlib.jl and here would surely help a lot.
E.g. it would be nice to fix FluxML/NNlib.jl#534

Verified fixed on current master (Julia 1.12.6, Metal.jl 1.10.0, GNNlib 1.3.0).

The scalar-indexing error came from propagate(copy_xj, g, +, …) calling adjacency_matrix → to_sparse → SparseArrays.sparse, which Metal cannot build without scalar indexing. #619 added a Metal-specific propagate(copy_xj, ::GNNGraph{<:METAL_COO_T}, …) override that uses a dense adjacency matmul (adjacency_matrix(g, …, fmt=:dense), built via NNlib.scatter), avoiding sparse entirely.

Re-running the documentation example (GCNConv/BatchNorm/GlobalPool/Dense GNNChain) on a Metal GPU now completes the full training step — forward pass, gradient, and Flux.update! — with the output an MtlArray and no scalar-indexing error. Closing.

Navigation Menu

Repository navigation

Documentation Example Causes Scalar Indexing Error wit Metal #618

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Documentation Example Causes Scalar Indexing Error wit Metal #618

Description

Activity

CarloLucibello commented on Sep 21, 2025

alonsoC1s commented on Nov 11, 2025

CarloLucibello commented on Dec 23, 2025

CarloLucibello commented on Jul 22, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions