From d68e71c4755d279cd6ed8cc4505ed01781951b01 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Mon, 29 Jun 2026 14:34:57 -0400 Subject: [PATCH] Add FunctionProperties extension: look through SciMLFunction wrappers `hasbranching` (FunctionProperties.jl) is used by SciMLSensitivity to decide whether a ReverseDiff tape can be compiled for an ODE right-hand side. It is called on `unwrapped_f(prob.f)`, i.e. an `AbstractSciMLFunction`, and analyzes whatever callable it is handed. Handed the SciMLFunction functor, it scans dispatch plumbing rather than the user's RHS: the operator-selection branch (`if f.f isa AbstractSciMLOperator`) and, for an MTK-compiled RHS, splat forwarding into a generated function. The result is that even a branch-free linear RHS reports `true`, and (without the companion FunctionProperties changes) genuine branches behind the forwarder are missed. This weakdep extension defines FunctionProperties.hasbranching(f::AbstractSciMLFunction, args...) = FunctionProperties.hasbranching(f.f, args...) so the analysis targets the wrapped RHS `f.f`. Loaded only when both packages are present; no new hard dependency. Stacked on SciML/FunctionProperties.jl#62 (the value-independent-branch detection that the unwrapped analysis relies on); compat set to that release. Co-Authored-By: Chris Rackauckas --- Project.toml | 8 ++++++-- ext/SciMLBaseFunctionPropertiesExt.jl | 16 ++++++++++++++++ test/function_properties_ext.jl | 15 +++++++++++++++ test/runtests.jl | 3 +++ 4 files changed, 40 insertions(+), 2 deletions(-) create mode 100644 ext/SciMLBaseFunctionPropertiesExt.jl create mode 100644 test/function_properties_ext.jl diff --git a/Project.toml b/Project.toml index cde7b602f..e27aac3a8 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "SciMLBase" uuid = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -version = "3.30.0" +version = "3.30.1" authors = ["Chris Rackauckas and contributors"] [deps] @@ -42,6 +42,7 @@ DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc" Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a" Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7" MonteCarloMeasurements = "0987c9cc-fe09-11e8-30f0-b96dd679fdca" @@ -61,6 +62,7 @@ SciMLBaseDifferentiationInterfaceExt = "DifferentiationInterface" SciMLBaseDistributionsExt = "Distributions" SciMLBaseEnzymeExt = "Enzyme" SciMLBaseForwardDiffExt = "ForwardDiff" +SciMLBaseFunctionPropertiesExt = "FunctionProperties" SciMLBaseMakieExt = "Makie" SciMLBaseMeasurementsExt = "Measurements" SciMLBaseMonteCarloMeasurementsExt = "MonteCarloMeasurements" @@ -90,6 +92,7 @@ DocStringExtensions = "0.9.5" EnumX = "1" Enzyme = "0.13" ForwardDiff = "0.10.36, 1" +FunctionProperties = "0.1.7" FunctionWrappersWrappers = "1" IteratorInterfaceExtensions = "^1" LinearAlgebra = "1.10" @@ -134,6 +137,7 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +FunctionProperties = "f62d2435-5019-4c03-9749-2d4c77af0cbc" PartialFunctions = "570af359-4316-4cb7-8c74-252c00c2016b" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" RCall = "6f49c342-dc21-5d91-9882-a32aef131414" @@ -148,4 +152,4 @@ UnicodePlots = "b8865327-cd53-5732-bb35-84acbb429228" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Aqua", "DifferentiationInterface", "Enzyme", "ForwardDiff", "PartialFunctions", "Pkg", "SafeTestsets", "SciMLTesting", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote"] +test = ["Aqua", "DifferentiationInterface", "Enzyme", "ForwardDiff", "FunctionProperties", "PartialFunctions", "Pkg", "SafeTestsets", "SciMLTesting", "Serialization", "StableRNGs", "StaticArrays", "Tables", "Test", "UnicodePlots", "Zygote"] diff --git a/ext/SciMLBaseFunctionPropertiesExt.jl b/ext/SciMLBaseFunctionPropertiesExt.jl new file mode 100644 index 000000000..c274ff8de --- /dev/null +++ b/ext/SciMLBaseFunctionPropertiesExt.jl @@ -0,0 +1,16 @@ +module SciMLBaseFunctionPropertiesExt + +using SciMLBase: AbstractSciMLFunction +using FunctionProperties: FunctionProperties + +# Analyze the wrapped right-hand side, not the `AbstractSciMLFunction` functor itself. The functor +# is dispatch plumbing: operator/`isa` selection (`if f.f isa AbstractSciMLOperator`) and, for an +# MTK-compiled RHS, splat forwarding into a generated function. Its branches are value-independent, +# and one of them (the dead operator path) is live IR that the scan would otherwise follow into +# arity-mismatch handling. Looking through to `f.f` is what makes `hasbranching` report the user's +# right-hand side, which is what callers (e.g. SciMLSensitivity, deciding whether a ReverseDiff +# tape can be compiled) actually care about. +FunctionProperties.hasbranching(f::AbstractSciMLFunction, args...) = + FunctionProperties.hasbranching(f.f, args...) + +end diff --git a/test/function_properties_ext.jl b/test/function_properties_ext.jl new file mode 100644 index 000000000..0d3cbf6cf --- /dev/null +++ b/test/function_properties_ext.jl @@ -0,0 +1,15 @@ +using SciMLBase, FunctionProperties, Test + +# The SciMLBaseFunctionPropertiesExt extension makes `hasbranching` look through an +# `AbstractSciMLFunction` to its wrapped right-hand side `f.f`, rather than analyzing the functor +# (whose operator-dispatch/forwarding branches are value-independent plumbing). +@test hasmethod(FunctionProperties.hasbranching, Tuple{SciMLBase.AbstractSciMLFunction, Vararg}) + +free!(du, u, p, t) = (du[1] = u[1] * u[1]; nothing) +branchy!(du, u, p, t) = (du[1] = u[1] > 0 ? u[1] : -u[1]; nothing) +u = [1.0] +du = [0.0] +p = SciMLBase.NullParameters() + +@test !hasbranching(ODEFunction(free!), copy(du), u, p, 0.0) +@test hasbranching(ODEFunction(branchy!), copy(du), u, p, 0.0) diff --git a/test/runtests.jl b/test/runtests.jl index f441739ba..4624f6886 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -25,6 +25,9 @@ run_tests(; @time @safetestset "Display" begin include("display.jl") end + @time @safetestset "FunctionProperties extension" begin + include("function_properties_ext.jl") + end @time @safetestset "Existence functions" begin include("existence_functions.jl") end