Generated function with an inlined linear SCC cannot be evaluated with SparseConnectivityTracer numbers

The failure is in the inlined linear solve: safe_ldiv calls numeric_ldiv! in ModelingToolkitTearing, which goes through LinearSolve and generic_lufact!, where the pivot comparison fails on tracer numbers. Start by running the reproducer, which traces the generated ODE function with TracerSparsityDetector and compares the result to the ForwardDiff Jacobian. Done means the traced pattern matches the nonzero pattern without the type-piracy overload, and the design question about the intended approach is answered by maintainers first.

Goal

We detect the sparsity pattern of the Jacobian of a compiled ModelingToolkit model with respect to its state, inputs and tunable parameters by evaluating the generated ODE function with SparseConnectivityTracer (TracerSparsityDetector, directly or through DifferentiationInterface). The alternative, Symbolics.jacobian_sparsity on the equations, traverses the expression trees without sharing common subexpressions; on MultibodyComponents models its cost grows exponentially with the number of bodies, so tracing the generated function is the only practical method for these models.

MultibodyComponents compiles its models with DefaultReassembleAlgorithm(; inline_linear_sccs = true, analytical_linear_scc_limit = 1), so the generated function contains inlined numeric linear solves.

Observed behaviour

With inline_linear_sccs = true, a linear SCC larger than analytical_linear_scc_limit is emitted as safe_ldiv(A, b), which calls ModelingToolkitTearing.numeric_ldiv! and solves the system with a LinearSolve LinearCache. When the generated function is called with tracer numbers, the LU factorization (generic_lufact! with RowMaximum pivoting) compares tracer values in the pivot search and throws

TypeError: non-boolean (SparseConnectivityTracer.GradientTracer{Int64, BitSet}) used in boolean context

A tracer does not carry a value, so the pivot comparison cannot be evaluated. SparseConnectivityTracer itself defines \ for matrices of tracers (every entry of the solution depends on every entry of A and b), but that method is not reached, because the solve goes through LinearSolve.

What we currently do

Downstream (JuliaComputing/MPCComponents.jl) we add a method to the internal function:

function ModelingToolkit.MTKTearing.numeric_ldiv!(A::AbstractMatrix{T}, b::AbstractVector{T}) where {T <: SparseConnectivityTracer.AbstractTracer}
    return copyto!(b, A \ b)
end

This is type piracy on a non-public function, and it relies on the current convention that numeric_ldiv! writes the solution into b (which JuliaComputing/StateSelection.jl#168 would change). With this method the traced pattern equals the nonzero pattern of the ForwardDiff Jacobian in the example below.

Question

Is there an intended way to evaluate a generated function that contains an inlined linear solve with SparseConnectivityTracer numbers (or with other number types for which the LU pivot search is not defined), short of compiling the model with inline_linear_sccs = false, which changes the compiled model and is not an option for MultibodyComponents models?

Reproducer

Plain ModelingToolkit, no MultibodyComponents. The algebraic unknowns y1, y2 form a 2x2 linear system whose coefficients depend on the state; analytical_linear_scc_limit = 1 makes it a numeric solve rather than an analytical one.

using ModelingToolkit, SparseConnectivityTracer, ForwardDiff
using ModelingToolkit: t_nounits as t, D_nounits as D

# The algebraic unknowns y1, y2 form a linear system whose coefficients depend on the state.
@variables x1(t) = 0.1 x2(t) = 0.2 x3(t) = 0.3 y1(t) y2(t)
eqs = [D(x1) ~ y1, D(x2) ~ y2, D(x3) ~ -x3,
       (2 + x1^2) * y1 + x2 * y2 ~ -x1,
       x1 * y1 + (3 + x2^2) * y2 ~ -x2]
@named model = System(eqs, t)
reassemble_alg = ModelingToolkit.StructuralTransformations.DefaultReassembleAlgorithm(;
    inline_linear_sccs = true, analytical_linear_scc_limit = 1)
sys = mtkcompile(model; reassemble_alg)
@show unknowns(sys)
prob = ODEProblem(sys, [], (0.0, 1.0))

f = (du, u) -> prob.f(du, u, prob.p, 0.0)
u = [0.4, -0.3, 0.2]
n = length(u)
J = ForwardDiff.jacobian(f, zeros(n), u)
println("ForwardDiff Jacobian:"); display(J)
println("Nonzeros of the ForwardDiff Jacobian:"); display(J .!= 0)

detector = TracerSparsityDetector()
try
    P = jacobian_sparsity(f, zeros(n), u, detector)
    println("Traced pattern:"); display(P)
catch err
    println("Tracing failed:")
    showerror(stdout, err, catch_backtrace()); println()
end

# Workaround currently used downstream (type piracy):
function ModelingToolkit.MTKTearing.numeric_ldiv!(A::AbstractMatrix{T}, b::AbstractVector{T}) where {T <: SparseConnectivityTracer.AbstractTracer}
    return copyto!(b, A \ b)
end
P = jacobian_sparsity(f, zeros(n), u, detector)
println("Traced pattern with the overload of numeric_ldiv!:"); display(P)
@show Matrix(P) == (J .!= 0)

Output

unknowns(sys) = SymbolicUtils.BasicSymbolicImpl.var"typeof(BasicSymbolicImpl)"{SymReal}[x3(t), x2(t), x1(t)]
ForwardDiff Jacobian:
3×3 Matrix{Float64}:
 -1.0  -0.0        -0.0
 -0.0  -0.297646    0.0569842
 -0.0  -0.0939958  -0.465547
Nonzeros of the ForwardDiff Jacobian:
3×3 BitMatrix:
 1  0  0
 0  1  1
 0  1  1
Tracing failed:
TypeError: non-boolean (SparseConnectivityTracer.GradientTracer{Int64, BitSet}) used in boolean context
Stacktrace:
  [1] generic_lufact!(A::Matrix{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, pivot::LinearAlgebra.RowMaximum; check::Bool, allowsingular::Bool)
  [5] lu(::Matrix{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}; kwargs::@Kwargs{check::Bool})
  [7] solve!(cache::LinearSolve.LinearCache{Matrix{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, Vector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, Vector{SparseConnectivityTra
 [10] solve!(::LinearSolve.LinearCache{Matrix{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, Vector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, Vector{SparseConnectivityTracer.G
 [14] solve_into!(cache::LinearSolve.LinearCache{Matrix{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, Vector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, Vector{SparseConnectivi
 [15] numeric_ldiv!(A::Matrix{SparseConnectivityTracer.GradientTracer{Int64, BitSet}}, b::Vector{SparseConnectivityTracer.GradientTracer{Int64, BitSet}})
 [16] safe_ldiv
 [19] generated_callfunc(::RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:__argₛᵧₘ8385217700582483117, :___mtkunknowns___, :___mtkparameters___, :__argₛᵧₘ11094091058616676188), Model
 [22] (::ODEFunction{true, SciMLBase.AutoDespecialize, ModelingToolkitBase.GeneratedFunctionWrapper{Tuple{2, 3, true}, RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:___mtkunknowns___, :___mtkpa
  ...
Traced pattern with the overload of numeric_ldiv!:
3×3 SparseArrays.SparseMatrixCSC{Bool, Int64} with 5 stored entries:
 1  ⋅  ⋅
 ⋅  1  1
 ⋅  1  1
Matrix(P) == (J .!= 0) = true

(Stack trace abridged; the relevant frames are generic_lufact! in LinearAlgebra, solve! in LinearSolve, solve_into!/numeric_ldiv!/safe_ldiv in ModelingToolkitTearing, called from the generated function.) The first lines of output also contain a PreallocationTools warning about the DiffCache size, caused by the ForwardDiff call.

Versions

Julia 1.12.7

Status `Project.toml`
  [a0c0ee7d] DifferentiationInterface v0.7.21
  [f6369f11] ForwardDiff v1.4.7
  [961ee093] ModelingToolkit v11.46.1
  [bbf590c4] OrdinaryDiffEqCore v4.19.1
  [53ae85a6] SciMLStructures v1.10.6
  [9f842d2f] SparseConnectivityTracer v1.2.4
  [2efcf032] SymbolicIndexingInterface v0.3.56

Also in the manifest: ModelingToolkitTearing 1.20.7, ModelingToolkitBase 1.79.0, LinearSolve 5.18.3, Symbolics 7.48.0, SymbolicUtils 4.55.0, SciMLBase 3.58.0. numeric_ldiv! is unchanged on the current default branch of JuliaComputing/StateSelection.jl.

Host: demeter2 · Claude Code session 20f98c68-57f3-4c7d-a6bb-90ad7a832832 · working directory: ~/.julia/dev/MPCComponents