Scalar Indexing Error in PINN Implementation on GPU

Dear All,

I am attempting to run the following Physics-Informed Neural Network (PINN) problem on a GPU. However, I encounter an error related to scalar indexing. Despite my efforts to debug this issue and explore potential solutions, I have been unable to find any relevant threads or guidance.
Could you please provide some advice or suggestions on how to address and resolve this scalar indexing error?

Thank you for your assistance.

code ===>

begin
    using ModelingToolkit, DomainSets
    using Lux
    using NeuralPDE
    using LuxCUDA, ComponentArrays, Random
end


# Parameters, variables, and derivatives
@parameters Xꜛ Yꜛₛₒₗ Yꜛₐ
@variables Θꜛₐ(..) ωꜛₐ(..) 
@variables Θꜛₛₒₗ(..) ξꜛₛₒₗ(..) 

DXꜛ = Differential(Xꜛ)
DYꜛₐ = Differential(Yꜛₐ)
DYꜛₛₒₗ = Differential(Yꜛₛₒₗ)
DYꜛₐ² = Differential(Yꜛₐ)^2
DYꜛₛₒₗ² = Differential(Yꜛₛₒₗ)^2

function Uꜛₛₒₗ(Yꜛₛₒₗ)
    y = Yꜛₛₒₗ * 1.2
    uₛₒₗ = 2 * y * (1.0 - 0.5 * y) / 4.0
    return uₛₒₗ / 1.5
end

function Uꜛₐ(Yꜛₐ)
    y = Yꜛₐ * 1.2
    uₐ = -1.2 - 5.0 *(3.0 ^ 2 - y ^ 2)
    return -uₐ / 2.0
end

i_fg(t_) = 4t_^3 - 3t_^4 + 10.0

function c2(Θ)
    t_ = Θ * 5.0 + 10.0
    return 10.0 * i_fg(t_) * 0.5 / (10.0 * 5000.0)
end
@register_symbolic c2(Θ)

# Equations
eqs = [
    Uꜛₐ(Yꜛₐ) * DXꜛ(Θꜛₐ(Xꜛ,Yꜛₐ)) ~ 5.0 * DYꜛₐ²(Θꜛₐ(Xꜛ,Yꜛₐ)),
    Uꜛₐ(Yꜛₐ) * DXꜛ(ωꜛₐ(Xꜛ,Yꜛₐ)) ~ 6.0 * DYꜛₐ²(ωꜛₐ(Xꜛ,Yꜛₐ)),
    Uꜛₛₒₗ(Yꜛₛₒₗ) * DXꜛ(Θꜛₛₒₗ(Xꜛ,Yꜛₛₒₗ)) ~ 8.0 * DYꜛₛₒₗ²(Θꜛₛₒₗ(Xꜛ,Yꜛₛₒₗ)),
    Uꜛₛₒₗ(Yꜛₛₒₗ) * DXꜛ(ξꜛₛₒₗ(Xꜛ,Yꜛₛₒₗ)) ~ 11.0 * DYꜛₛₒₗ²(ξꜛₛₒₗ(Xꜛ,Yꜛₛₒₗ))
]

bcs = [
    Θꜛₐ(1.0, Yꜛₐ) ~ 1.0,
    ωꜛₐ(1.0, Yꜛₐ) ~ 1.0,
    Θꜛₛₒₗ(0.0, Yꜛₛₒₗ) ~ 1.0,
    ξꜛₛₒₗ(0.0, Yꜛₛₒₗ) ~ 1.0,

    DYꜛₐ(Θꜛₐ(Xꜛ, 0.0)) ~ 0.0,
    DYꜛₐ(ωꜛₐ(Xꜛ, 0.0)) ~ 0.0,
    Θꜛₛₒₗ(Xꜛ, 0.0) ~ 0.0,
    DYꜛₛₒₗ(ξꜛₛₒₗ(Xꜛ, 0.0)) ~ 0.0,

    DYꜛₐ(Θꜛₐ(Xꜛ, 1.0)) + 5.0 * DYꜛₛₒₗ(Θꜛₛₒₗ(Xꜛ, 1.0)) - 5.0 * c2(Θꜛₛₒₗ(Xꜛ, 1.0)) * DYꜛₛₒₗ(ξꜛₛₒₗ(Xꜛ, 1.0)) ~ 0.0,
    DYꜛₐ(ωꜛₐ(Xꜛ, 1.0)) + 1.05 * DYꜛₛₒₗ(ξꜛₛₒₗ(Xꜛ, 1.0)) ~ 0.0
]

domains = [
    Xꜛ ∈ Interval(0.0, 1.0),
    Yꜛₐ ∈ Interval(0.0, 1.0),
    Yꜛₛₒₗ ∈ Interval(0.0, 1.0),
]

@named pdesystem = PDESystem(eqs, bcs, domains,[Xꜛ, Yꜛₐ, Yꜛₛₒₗ], [Θꜛₐ(Xꜛ,Yꜛₐ), ωꜛₐ(Xꜛ,Yꜛₐ), Θꜛₛₒₗ(Xꜛ,Yꜛₛₒₗ), ξꜛₛₒₗ(Xꜛ,Yꜛₛₒₗ)])

# Neural network
dim = 2 # number of dimensions
N_nr = 8 # number of neurons
chain = [Lux.Chain(Dense(dim, N_nr, Lux.σ), Dense(N_nr, N_nr, Lux.σ), Dense(N_nr, 1)) for _ in 1:4]

strategy = QuasiRandomTraining(100)
gpu_dev = gpu_device()
# Seeding
rng = Random.default_rng()
Random.seed!(rng, 0)

ps = Lux.setup(Random.default_rng(), chain)[1] 
ps = ps |> gpu_dev
ps = f64(ps)


discretization = PhysicsInformedNN(chain, strategy,init_params = ps)
prob = discretize(pdesystem, discretization)

error ===>

ERROR: 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. 

I think this issue arises from c2(Θꜛₛₒₗ(Xꜛ, 1.0)) in Bcs but I don’t know how to handle it. It might stem form scalar indexing in discretization.

thanks in advance for your assistance and taking the time to help.

@ChrisRackauckas

error ===>

ERROR: 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 C:\Users\lab30501\.julia\packages\GPUArraysCore\GMsgk\src\GPUArraysCore.jl:155        
  [3] _assertscalar(op::String, behavior::GPUArraysCore.ScalarIndexing)
    @ GPUArraysCore C:\Users\lab30501\.julia\packages\GPUArraysCore\GMsgk\src\GPUArraysCore.jl:128        
  [4] assertscalar(op::String)
    @ GPUArraysCore C:\Users\lab30501\.julia\packages\GPUArraysCore\GMsgk\src\GPUArraysCore.jl:116        
  [5] getindex(A::CuArray{Float64, 2, CUDA.DeviceMemory}, I::Int64)
    @ GPUArrays C:\Users\lab30501\.julia\packages\GPUArrays\8Y80U\src\host\indexing.jl:48
  [6] iterate
    @ .\abstractarray.jl:1217 [inlined]
  [7] iterate
    @ .\abstractarray.jl:1215 [inlined]
  [8] _zip_iterate_some
    @ .\iterators.jl:428 [inlined]
  [9] _zip_iterate_some
    @ .\iterators.jl:430 [inlined]
 [10] _zip_iterate_all
    @ .\iterators.jl:420 [inlined]
 [11] iterate
    @ .\iterators.jl:410 [inlined]
 [12] _append!(a::Vector{Float64}, ::Base.HasShape{2}, iter::CuArray{Float64, 2, CUDA.DeviceMemory})      
    @ Base .\array.jl:1197
 [13] append!
    @ .\array.jl:1187 [inlined]
 [14] make_idx(data::Vector{Float64}, x::CuArray{Float64, 2, CUDA.DeviceMemory}, last_val::Int64)
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:186   
 [15] #28
    @ C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:161 [inlined]
 [16] (::ComponentArrays.var"#28#29"{Vector{…}, Base.RefValue{…}})(::Pair{Symbol, CuArray{…}})
    @ ComponentArrays .\none:0
 [17] iterate
    @ .\generator.jl:47 [inlined]
 [18] merge(a::@NamedTuple{}, itr::Base.Generator{@Kwargs{…}, ComponentArrays.var"#28#29"{…}})
    @ Base .\namedtuple.jl:361
 [19] make_idx(data::Vector{Float64}, nt::@NamedTuple{weight::CuArray{…}, bias::CuArray{…}}, last_val::Int64)
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:159   
 [20] #28
    @ C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:161 [inlined]
 [21] (::ComponentArrays.var"#28#29"{Vector{…}, Base.RefValue{…}})(::Pair{Symbol, @NamedTuple{…}})
    @ ComponentArrays .\none:0
 [22] iterate
    @ .\generator.jl:47 [inlined]
 [23] merge(a::@NamedTuple{}, itr::Base.Generator{@Kwargs{…}, ComponentArrays.var"#28#29"{…}})
    @ Base .\namedtuple.jl:361
 [24] make_idx(data::Vector{…}, nt::@NamedTuple{…}, last_val::Int64)
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:159   
 [25] #28
    @ C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:161 [inlined]
 [26] (::ComponentArrays.var"#28#29"{Vector{…}, Base.RefValue{…}})(::Pair{Symbol, @NamedTuple{…}})        
    @ ComponentArrays .\none:0
 [27] iterate
    @ .\generator.jl:47 [inlined]
 [28] merge(a::@NamedTuple{}, itr::Base.Generator{@Kwargs{…}, ComponentArrays.var"#28#29"{…}})
    @ Base .\namedtuple.jl:361
 [29] make_idx(data::Vector{…}, nt::@NamedTuple{…}, last_val::Int64)
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:159   
 [30] make_carray_args(A::Type{…}, nt::@NamedTuple{…})
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:151   
 [31] make_carray_args(nt::@NamedTuple{…})
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:144   
 [32] ComponentArray(nt::@NamedTuple{Θꜛₐ::@NamedTuple{…}, ωꜛₐ::@NamedTuple{…}, Θꜛₛₒₗ::@NamedTuple{…}, ξꜛₛₒₗ::@NamedTuple{…}})
    @ ComponentArrays C:\Users\lab30501\.julia\packages\ComponentArrays\joxVV\src\componentarray.jl:64    
 [33] symbolic_discretize(pde_system::PDESystem, discretization::PhysicsInformedNN{…})
    @ NeuralPDE C:\Users\lab30501\.julia\packages\NeuralPDE\z18Qg\src\discretize.jl:415
 [34] discretize(pde_system::PDESystem, discretization::PhysicsInformedNN{…})
    @ NeuralPDE C:\Users\lab30501\.julia\packages\NeuralPDE\z18Qg\src\discretize.jl:708
 [35] top-level scope
    @ d:\Maysam\01 - Projects\05_Ionic_Liquid_Dsign_Simulation\_Julia_code\_NPDE_2d\MWE.jl:88
Some type information was truncated. Use `show(err)` to see complete types.

@SathvikBhagavan is the one that has been working on related issues. Is this known?

Thanks for your response.

I changed

to
ps = ps .|> ComponentArray |> gpu_dev.

It works.