Fix GPU scalar indexing in BatchExaModel structure queries

BatchExaModel.jac_structure! and hess_structure! now pass getbackend(m)
to backend-aware _jac_structure!/_obj_hess_structure!/_con_hess_structure!
so the KA extension can intercept and use GPU kernels instead of
scalar-indexing f.itr arrays.

Also converts obj_weight to model type T in batch hess_coord! to prevent
Metal InvalidIRError from Float64 values on Float32 models.

FlatNLPModel structure queries now use device-native temp arrays for the
per-instance query, then copy to CPU for replication.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: sshin23

src/BatchNLPModels.jl

@@ -226,12 +226,18 @@ function NLPModels.jac_nln_structure!(m::FlatNLPModel, rows::AbstractVector{<:In
     ncon = NLPModels.get_ncon(m.batch)
     nnzj = NLPModels.get_nnzj(m.batch)
 
-    # Compute structure on CPU (structure queries use scalar indexing)
+    # Use device-native arrays for the per-instance query (avoids scalar indexing on GPU itr)
+    r1_dev = similar(m.batch.meta.x0, Int, nnzj)
+    c1_dev = similar(m.batch.meta.x0, Int, nnzj)
+    NLPModels.jac_structure!(m.batch, r1_dev, c1_dev)
+
+    # Copy to CPU for replication
+    r1 = Vector{Int}(r1_dev)
+    c1 = Vector{Int}(c1_dev)
     r_cpu = Vector{Int}(undef, nnzj * nb)
     c_cpu = Vector{Int}(undef, nnzj * nb)
-    r1 = @view r_cpu[1:nnzj]
-    c1 = @view c_cpu[1:nnzj]
-    NLPModels.jac_structure!(m.batch, r1, c1)
+    copyto!(r_cpu, 1, r1, 1, nnzj)
+    copyto!(c_cpu, 1, c1, 1, nnzj)
 
     # Replicate for each instance with shifted indices
     @inbounds for s in 2:nb
@@ -261,12 +267,18 @@ function NLPModels.hess_structure!(m::FlatNLPModel, rows::AbstractVector{<:Integ
     nvar = NLPModels.get_nvar(m.batch)
     nnzh = NLPModels.get_nnzh(m.batch)
 
-    # Compute structure on CPU (structure queries use scalar indexing)
+    # Use device-native arrays for the per-instance query (avoids scalar indexing on GPU)
+    r1_dev = similar(m.batch.meta.x0, Int, nnzh)
+    c1_dev = similar(m.batch.meta.x0, Int, nnzh)
+    NLPModels.hess_structure!(m.batch, r1_dev, c1_dev)
+
+    # Copy to CPU for replication
+    r1 = Vector{Int}(r1_dev)
+    c1 = Vector{Int}(c1_dev)
     r_cpu = Vector{Int}(undef, nnzh * nb)
     c_cpu = Vector{Int}(undef, nnzh * nb)
-    r1 = @view r_cpu[1:nnzh]
-    c1 = @view c_cpu[1:nnzh]
-    NLPModels.hess_structure!(m.batch, r1, c1)
+    copyto!(r_cpu, 1, r1, 1, nnzh)
+    copyto!(c_cpu, 1, c1, 1, nnzh)
 
     # Replicate for each instance with shifted indices
     @inbounds for s in 2:nb

src/nlp.jl

@@ -1329,6 +1329,8 @@ _jac_structure!(T, cons::Tuple{}, rows, cols) = nothing
     _jac_structure!(T, Base.tail(cons), rows, cols)
     sjacobian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN))
 end
+# Backend-aware fallbacks (KA extension overrides these for GPU backends)
+_jac_structure!(T, ::Nothing, cons, rows, cols) = _jac_structure!(T, cons, rows, cols)
 
 function hess_structure!(m::AbstractExaModel{T}, rows::AbstractVector, cols::AbstractVector) where T
     _obj_hess_structure!(T, m.objs, rows, cols)
@@ -1341,12 +1343,16 @@ _obj_hess_structure!(T, objs::Tuple{}, rows, cols) = nothing
     _obj_hess_structure!(T, Base.tail(objs), rows, cols)
     shessian!(rows, cols, first(objs), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN))
 end
+# Backend-aware fallback (KA extension overrides for GPU backends)
+_obj_hess_structure!(T, ::Nothing, objs, rows, cols) = _obj_hess_structure!(T, objs, rows, cols)
 
 _con_hess_structure!(T, cons::Tuple{}, rows, cols) = nothing
 @inline function _con_hess_structure!(T, cons::Tuple, rows, cols)
     _con_hess_structure!(T, Base.tail(cons), rows, cols)
     shessian!(rows, cols, first(cons), NaNSource{T}(), NaNSource{T}(), T(NaN), T(NaN))
 end
+# Backend-aware fallback (KA extension overrides for GPU backends)
+_con_hess_structure!(T, ::Nothing, cons, rows, cols) = _con_hess_structure!(T, cons, rows, cols)
 
 # ============================================================================
 # Batch-aware evaluation — all low-level functions loop over 1:nb,
@@ -1890,7 +1896,7 @@ function NLPModels.jac_structure!(
     rows::AbstractVector{<:Integer},
     cols::AbstractVector{<:Integer},
 ) where {T}
-    _jac_structure!(T, m.cons, rows, cols)
+    _jac_structure!(T, getbackend(m), m.cons, rows, cols)
     return rows, cols
 end
 
@@ -1909,8 +1915,9 @@ function NLPModels.hess_structure!(
     rows::AbstractVector{<:Integer},
     cols::AbstractVector{<:Integer},
 ) where {T}
-    _obj_hess_structure!(T, m.objs, rows, cols)
-    _con_hess_structure!(T, m.cons, rows, cols)
+    backend = getbackend(m)
+    _obj_hess_structure!(T, backend, m.objs, rows, cols)
+    _con_hess_structure!(T, backend, m.cons, rows, cols)
     return rows, cols
 end
 
@@ -1928,7 +1935,7 @@ function NLPModels.hess_coord!(
     npar = Base.size(m.θ, 1)
     nnzh = NLPModels.get_nnzh(m)
     backend = getbackend(m)
-    _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight, nb, nvar, npar, nnzh, backend)
+    _obj_hess_coord!(m.objs, vec(bx), vec(m.θ), vec(hvals), obj_weight isa Number ? T(obj_weight) : obj_weight, nb, nvar, npar, nnzh, backend)
     _con_hess_coord!(m.cons, vec(bx), vec(m.θ), vec(by), vec(hvals), nb, nvar, npar, ncon, nnzh, backend)
     return hvals
 end