refactor: move ChainRulesCoreExt into main package

Author: AayushSabharwal

Project.toml

@@ -8,6 +8,7 @@ ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 BlockArrays = "8e7c35d0-a365-5155-bbbb-fb81a777f24e"
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa"
 CommonSolve = "38540f10-b2f7-11e9-35d8-d573e4eb0ff2"
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -65,7 +66,6 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 [weakdeps]
 BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 CasADi = "c49709b8-5c63-11e9-2fb2-69db5844192f"
-ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
@@ -74,7 +74,6 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 [extensions]
 MTKBifurcationKitExt = "BifurcationKit"
 MTKCasADiDynamicOptExt = "CasADi"
-MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDeepDiffsExt = "DeepDiffs"
 MTKFMIExt = "FMI"
 MTKInfiniteOptExt = "InfiniteOpt"

src/ModelingToolkit.jl

@@ -62,6 +62,8 @@ import BlockArrays: BlockArray, BlockedArray, Block, blocksize, blocksizes, bloc
 using OffsetArrays: Origin
 import CommonSolve
 import EnumX
+import ChainRulesCore
+import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
 
 using RuntimeGeneratedFunctions
 using RuntimeGeneratedFunctions: drop_expr
@@ -204,6 +206,8 @@ include("structural_transformation/StructuralTransformations.jl")
 @reexport using .StructuralTransformations
 include("inputoutput.jl")
 
+include("adjoints.jl")
+
 for S in subtypes(ModelingToolkit.AbstractSystem)
     S = nameof(S)
     @eval convert_system(::Type{<:$S}, sys::$S) = sys

ext/MTKChainRulesCoreExt.jl renamed to src/adjoints.jl

@@ -1,17 +1,11 @@
-module MTKChainRulesCoreExt
-
-import ModelingToolkit as MTK
-import ChainRulesCore
-import ChainRulesCore: Tangent, ZeroTangent, NoTangent, zero_tangent, unthunk
-
-function ChainRulesCore.rrule(::Type{MTK.MTKParameters}, tunables, args...)
+function ChainRulesCore.rrule(::Type{MTKParameters}, tunables, args...)
     function mtp_pullback(dt)
         dt = unthunk(dt)
         dtunables = dt isa AbstractArray ? dt : dt.tunable
         (NoTangent(), dtunables[1:length(tunables)],
             ntuple(_ -> NoTangent(), length(args))...)
     end
-    MTK.MTKParameters(tunables, args...), mtp_pullback
+    MTKParameters(tunables, args...), mtp_pullback
 end
 
 function subset_idxs(idxs, portion, template)
@@ -70,23 +64,23 @@ function selected_tangents(
 end
 
 function ChainRulesCore.rrule(
-        ::typeof(MTK.remake_buffer), indp, oldbuf::MTK.MTKParameters, idxs, vals)
+        ::typeof(remake_buffer), indp, oldbuf::MTKParameters, idxs, vals)
     if idxs isa AbstractSet
         idxs = collect(idxs)
     end
     idxs = map(idxs) do i
-        i isa MTK.ParameterIndex ? i : MTK.parameter_index(indp, i)
+        i isa ParameterIndex ? i : parameter_index(indp, i)
     end
-    newbuf = MTK.remake_buffer(indp, oldbuf, idxs, vals)
+    newbuf = remake_buffer(indp, oldbuf, idxs, vals)
     tunable_idxs = reduce(
-        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Tunable);
+        vcat, (idx.idx for idx in idxs if idx.portion isa SciMLStructures.Tunable);
         init = Union{Int, AbstractVector{Int}}[])
     initials_idxs = reduce(
-        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Initials);
+        vcat, (idx.idx for idx in idxs if idx.portion isa SciMLStructures.Initials);
         init = Union{Int, AbstractVector{Int}}[])
-    disc_idxs = subset_idxs(idxs, MTK.SciMLStructures.Discrete(), oldbuf.discrete)
-    const_idxs = subset_idxs(idxs, MTK.SciMLStructures.Constants(), oldbuf.constant)
-    nn_idxs = subset_idxs(idxs, MTK.NONNUMERIC_PORTION, oldbuf.nonnumeric)
+    disc_idxs = subset_idxs(idxs, SciMLStructures.Discrete(), oldbuf.discrete)
+    const_idxs = subset_idxs(idxs, SciMLStructures.Constants(), oldbuf.constant)
+    nn_idxs = subset_idxs(idxs, NONNUMERIC_PORTION, oldbuf.nonnumeric)
 
     pullback = let idxs = idxs
         function remake_buffer_pullback(buf′)
@@ -102,13 +96,11 @@ function ChainRulesCore.rrule(
             oldbuf′ = Tangent{typeof(oldbuf)}(;
                 tunable, initials, discrete, constant, nonnumeric)
             idxs′ = NoTangent()
-            vals′ = map(i -> MTK._ducktyped_parameter_values(buf′, i), idxs)
+            vals′ = map(i -> _ducktyped_parameter_values(buf′, i), idxs)
             return f′, indp′, oldbuf′, idxs′, vals′
         end
     end
     newbuf, pullback
 end
 
-ChainRulesCore.@non_differentiable Base.getproperty(sys::MTK.AbstractSystem, x::Symbol)
-
-end
+ChainRulesCore.@non_differentiable Base.getproperty(sys::AbstractSystem, x::Symbol)