MTKBase: FunctionProperties extension exempting MTKParameters indexing

lib/ModelingToolkitBase/Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkitBase"
 uuid = "7771a370-6774-4173-bd38-47e70ca0b839"
 authors = ["Yingbo Ma <mayingbo5@gmail.com>", "Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "1.48.1"
+version = "1.48.2"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -68,6 +68,7 @@ CasADi = "c49709b8-5c63-11e9-2fb2-69db5844192f"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DiffEqNoiseProcess = "77a26b50-5914-5dd7-bc55-306e6241c503"
 DynamicQuantities = "06fc5a27-2a28-4c7c-a15d-362465fb6821"
+FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc"
 InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
 JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
@@ -82,6 +83,7 @@ MTKCasADiDynamicOptExt = "CasADi"
 MTKChainRulesCoreExt = "ChainRulesCore"
 MTKDiffEqNoiseProcessExt = "DiffEqNoiseProcess"
 MTKDynamicQuantitiesExt = "DynamicQuantities"
+MTKFunctionPropertiesExt = "FunctionProperties"
 MTKInfiniteOptExt = "InfiniteOpt"
 MTKJuliaFormatterExt = "JuliaFormatter"
 MTKLabelledArraysExt = "LabelledArrays"
@@ -125,6 +127,7 @@ EnzymeCore = "0.8"
 ExprTools = "0.1.10"
 FillArrays = "1.13.0"
 ForwardDiff = "0.10.3, 1"
+FunctionProperties = "0.1.7"
 FunctionWrappers = "1.1.2"
 FunctionWrappersWrappers = "0.1, 1"
 Graphs = "1.5.2"

lib/ModelingToolkitBase/ext/MTKFunctionPropertiesExt.jl

@@ -0,0 +1,16 @@
+module MTKFunctionPropertiesExt
+
+import ModelingToolkitBase as MTKBase
+import FunctionProperties
+
+# Integer indexing into an `MTKParameters` buffer selects one of the parameter portions by index.
+# The selecting branch — `getindex(::MTKParameters, ::Int)` is a `@generated` `if idx == k` chain —
+# is value-independent: every real call site in generated right-hand-side code passes a literal
+# index that constant-folds the branch away, but `hasbranching`'s recursion only sees the widened
+# `Int`, so it reports the branch. Mark this plumbing as a leaf so `hasbranching` reflects the
+# user's RHS (e.g. a registered branching activation) rather than parameter-container indexing.
+FunctionProperties.is_leaf_sig(
+    ::Type{<:Tuple{typeof(Base.getindex), <:MTKBase.MTKParameters, Vararg}}
+) = true
+
+end

lib/ModelingToolkitBase/test/extensions/Project.toml

@@ -4,6 +4,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
 ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839"
@@ -26,6 +27,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 ModelingToolkitBase = {path = "../.."}
 
 [compat]
+FunctionProperties = "0.1.7"
 SafeTestsets = "0.1, 1"
 SciMLTesting = "1"
 TaylorDiff = "0.3.5"

lib/ModelingToolkitBase/test/extensions/function_properties.jl

@@ -0,0 +1,24 @@
+using ModelingToolkitBase, FunctionProperties, Test
+using ModelingToolkitBase: t_nounits as t, D_nounits as D
+
+# `getindex(::MTKParameters, ::Int)` selects a parameter buffer by index. Its branch is on the
+# index (constant at every real generated-code call site), never on values, so `hasbranching`
+# should not report it. The MTKFunctionPropertiesExt extension marks it a leaf via `is_leaf_sig`.
+
+@parameters a = 2.0
+@variables x(t) = 1.0
+@named sys = System([D(x) ~ a * x], t)
+sys = mtkcompile(sys)
+prob = ODEProblem(sys, [x => 1.0], (0.0, 1.0))
+p = prob.p
+
+@test p isa ModelingToolkitBase.MTKParameters
+
+# The extension's signature-level exemption is active for the parameter-container indexing...
+@test FunctionProperties.is_leaf_sig(Tuple{typeof(getindex), typeof(p), Int})
+# ...and does not over-reach to ordinary containers.
+@test !FunctionProperties.is_leaf_sig(Tuple{typeof(getindex), Vector{Float64}, Int})
+
+# A function that only indexes into the parameter container therefore has no value-dependent branch.
+index_param(pp) = @inbounds pp[1]
+@test !hasbranching(index_param, p)

lib/ModelingToolkitBase/test/runtests.jl

@@ -130,6 +130,7 @@ end
         activate_extensions_env()
         @safetestset "HomotopyContinuation Extension Test" include("extensions/homotopy_continuation.jl")
         @safetestset "LabelledArrays Test" include("extensions/labelledarrays.jl")
+        @safetestset "FunctionProperties Test" include("extensions/function_properties.jl")
         @safetestset "BifurcationKit Extension Test" include("extensions/bifurcationkit.jl")
         # @safetestset "Auto Differentiation Test" include("extensions/ad.jl")
     end