L2_NORM resorts to scalar indexing on GPU ComponentArray

[vpuri3]

julia> using CUDA, NonlinearSolve; CUDA.allowscalar(false)                                         
                                                 
julia> x = ComponentArray(; x = ones(2), y = ones(2)) |> cu              
ComponentVector{Float32, CuArray{Float32, 1, CUDA.DeviceMemory}, Tuple{Axis{(x = 1:2, y = 3:4)}}}(x = Float32[1.0, 1.0], y = Float32[1.0, 1.0])                                                       
                                                                                                   
julia> NonlinearSolveBase.L2_NORM(x)                                                               
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.

Need a ComponentArray extension with L2_NORM(x::CA) = L2_NORM(getdata(x)).

[avik-pal]

does NonlinearSolve work with a CA initial guess (even on CPU)?

[vpuri3]

I've been using NonlinearSolve.jl to optimize neural networks (written in Lux with ComponentArrays) here. Haven't had an issue before, though this is the first time I'm testing it since the reorg.

[vpuri3]

This solved the issue for me. https://github.com/vpuri3/NeuralROMs.jl/blob/master/experiments_SNFROM/compare.jl#L28-L30

# https://github.com/SciML/NonlinearSolve.jl/issues/524                                       
using NonlinearSolve, ComponentArrays                                                         
NonlinearSolveBase.L2_NORM(u::ComponentArray) = NonlinearSolveBase.L2_NORM(getdata(u))