Merge pull request #3100 from DhairyaLGandhi/dg/crc3

Allow arrays in `MTKParameters` pullback
SciML · Oct 7, 2024 · 284e463 · 284e463
2 parents d7fa2b9 + 91e5e70
commit 284e463
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ext/MTKChainRulesCoreExt.jl b/ext/MTKChainRulesCoreExt.jl
@@ -7,7 +7,8 @@ import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
 function ChainRulesCore.rrule(::Type{MTK.MTKParameters}, tunables, args...)
     function mtp_pullback(dt)
         dt = unthunk(dt)
-        (NoTangent(), dt.tunable[1:length(tunables)],
+        dtunables = dt isa AbstractArray ? dt : dt.tunable
+        (NoTangent(), dtunables[1:length(tunables)],
             ntuple(_ -> NoTangent(), length(args))...)
     end
     MTK.MTKParameters(tunables, args...), mtp_pullback