Adapt setup for DiffEqGPU DAEs

Author: ChrisRackauckas

src/adapt.jl

@@ -16,9 +16,25 @@ function adapt_structure(to,
         adapt(to, prob.kwargs)...)
 end
 
-function adapt_structure(to, f::ODEFunction{iip}) where {iip}
-    if f.mass_matrix !== I && f.initialization_data !== nothing
-        error("Adaptation to GPU failed: DAEs of ModelingToolkit currently not supported.")
-    end
-    ODEFunction{iip, FullSpecialize}(f.f, jac = f.jac, mass_matrix = f.mass_matrix)
+# Allow DAE adaptation for GPU kernels
+function adapt_structure(to, f::SciMLBase.ODEFunction{iip}) where {iip}
+    # For GPU kernels, we now support DAEs with mass matrices and initialization
+    SciMLBase.ODEFunction{iip, SciMLBase.FullSpecialize}(
+        f.f,
+        jac = f.jac,
+        mass_matrix = f.mass_matrix,
+        initialization_data = f.initialization_data
+    )
+end
+
+# Adapt OverrideInitData for GPU compatibility
+function adapt_structure(to, f::SciMLBase.OverrideInitData)
+    SciMLBase.OverrideInitData(
+        adapt(to, f.initializeprob),  # Also adapt initializeprob
+        f.update_initializeprob!,
+        f.initializeprobmap,
+        f.initializeprobpmap,
+        nothing,  # Set metadata to nothing for GPU compatibility
+        f.is_update_oop
+    )
 end