Merge pull request #780 from AayushSabharwal/as/depwarn

fix: fix depwarn in `responsible_map`
SciML · Sep 4, 2024 · 36e0c3a · 36e0c3a
2 parents 8763f7b + 7336f2d
commit 36e0c3a
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/src/ensemble/basic_ensemble_solve.jl b/src/ensemble/basic_ensemble_solve.jl
@@ -233,12 +233,16 @@ function solve_batch(prob, alg, ensemblealg::EnsembleDistributed, II, pmap_batch
     tighten_container_eltype(batch_data)
 end
 
+__getindex(x, i) = x[i]
+__getindex(x::AbstractVectorOfArray, i) = x.u[i]
+
 function responsible_map(f, II...)
-    batch_data = Vector{Core.Compiler.return_type(f, Tuple{typeof.(getindex.(II, 1))...})}(
+    batch_data = Vector{Core.Compiler.return_type(
+        f, Tuple{ntuple(i -> typeof(__getindex(II[i], 1)), Val(length(II)))...})}(
         undef,
         length(II[1]))
     for i in 1:length(II[1])
-        batch_data[i] = f(getindex.(II, i)...)
+        batch_data[i] = f(ntuple(ii -> __getindex(II[ii], i), Val(length(II)))...)
     end
     batch_data
 end