Differentiate through the values in `p`

[ayushinav]

Describe the bug 🐞
When sampling using NUTS() to sample the parameters of the equation to be solved p, the following error is returned

Float64(::ForwardDiff.Dual{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag, Float64}, Float64, 1})

Additional context

For a type stable function f, if any of the values in p is ForwardDiff.Dual type, https://github.com/SciML/NonlinearSolve.jl/blob/master/lib/BracketingNonlinearSolve/src/itp.jl#L111 returns a δ that isFloat64 type, which then raises the above error in https://github.com/SciML/NonlinearSolve.jl/blob/master/lib/BracketingNonlinearSolve/src/itp.jl#L114 because diff will be a Dual number.

Minimal Reproducible Example

p = (; a = 0.1)

prob_init = IntervalNonlinearProblem(
            f_test, (1.0f-15, 1.0f0), p)
sol = solve(prob_init)

solve(prob_init; p = (; a = 0.3))
return sol.u

@model function prob_test(u_true, u_err)
    a_ ~ truncated(Normal(0.5, 0.25); lower = 0.1, upper = 1.)
    p = (; a = a_)

    prob_init = IntervalNonlinearProblem(
            f_test, (1.0f-15, 1.0f0), p)
    sol = solve(prob_init)
    # @show sol.u

    u_true ~ Normal(sol.u, u_err)

end

u_t = 0.3
u_err = 1e-6

model_test = prob_test(u_t, u_err)

chain_test = sample(model_test, Prior(), 1000)

chain_test = sample(model_test, NUTS(), 1000)

Error & Stacktrace ⚠️

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag, Float64}, Float64, 1})

Closest candidates are:
  (::Type{T})(::Real, ::RoundingMode) where T<:AbstractFloat
   @ Base rounding.jl:207
  (::Type{T})(::T) where T<:Number
   @ Core boot.jl:792
  Float64(::UInt8)
   @ Base float.jl:165
  ...

Stacktrace:
  [1] copysign(x::Float64, y::ForwardDiff.Dual{ForwardDiff.Tag{DynamicPPL.DynamicPPLTag, Float64}, Float64, 1})
    @ Base ./floatfuncs.jl:8
  [2] __solve(::IntervalNonlinearProblem{…}, ::ITP{…}; maxiters::Int64, abstol::Nothing, verbose::Bool, kwargs::@Kwargs{})
    @ BracketingNonlinearSolve ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/itp.jl:114
  [3] __solve
    @ ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/itp.jl:59 [inlined]
  [4] #bracketingnonlinear_solve_up#13
    @ ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/BracketingNonlinearSolve.jl:39 [inlined]
  [5] bracketingnonlinear_solve_up
    @ ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/BracketingNonlinearSolve.jl:38 [inlined]
  [6] #solve#12
    @ ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/BracketingNonlinearSolve.jl:34 [inlined]
  [7] solve
    @ ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/BracketingNonlinearSolve.jl:32 [inlined]
  [8] #solve#10
    @ ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/BracketingNonlinearSolve.jl:25 [inlined]
  [9] solve(prob::IntervalNonlinearProblem{…})
    @ BracketingNonlinearSolve ~/.julia/packages/BracketingNonlinearSolve/KYxGw/src/BracketingNonlinearSolve.jl:24
 [10] prob_test(__model__::DynamicPPL.Model{…}, __varinfo__::DynamicPPL.TypedVarInfo{…}, __context__::DynamicPPL.SamplingContext{…}, u_true::Float64, u_err::Float64)
    @ Main ~/Desktop/pkgs/wip/kivu_rp.jl:184
 [11] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/f9pFn/src/model.jl:936 [inlined]
 [12] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/f9pFn/src/model.jl:909 [inlined]
 [13] evaluate!!
    @ ~/.julia/packages/DynamicPPL/f9pFn/src/model.jl:857 [inlined]
 [14] logdensity_at(x::Vector{…}, model::DynamicPPL.Model{…}, varinfo::DynamicPPL.TypedVarInfo{…}, context::DynamicPPL.SamplingContext{…})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/f9pFn/src/logdensityfunction.jl:184
 [15] FixTail
    @ ~/.julia/packages/DifferentiationInterface/zJHX8/src/utils/context.jl:169 [inlined]
 [16] vector_mode_dual_eval!
    @ ~/.julia/packages/ForwardDiff/UBbGT/src/apiutils.jl:24 [inlined]
 [17] vector_mode_gradient!
    @ ~/.julia/packages/ForwardDiff/UBbGT/src/gradient.jl:98 [inlined]
 [18] gradient!(result::DiffResults.MutableDiffResult{…}, f::DifferentiationInterface.FixTail{…}, x::Vector{…}, cfg::ForwardDiff.GradientConfig{…}, ::Val{…})
    @ ForwardDiff ~/.julia/packages/ForwardDiff/UBbGT/src/gradient.jl:39
 [19] value_and_gradient(::typeof(DynamicPPL.logdensity_at), ::DifferentiationInterfaceForwardDiffExt.ForwardDiffGradientPrep{…}, ::AutoForwardDiff{…}, ::Vector{…}, ::DifferentiationInterface.Constant{…}, ::DifferentiationInterface.Constant{…}, ::DifferentiationInterface.Constant{…})
    @ DifferentiationInterfaceForwardDiffExt ~/.julia/packages/DifferentiationInterface/zJHX8/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl:417
 [20] logdensity_and_gradient(f::LogDensityFunction{…}, x::Vector{…})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/f9pFn/src/logdensityfunction.jl:215
 [21] Fix1
    @ ./operators.jl:1118 [inlined]
 [22] ∂H∂θ(h::AdvancedHMC.Hamiltonian{…}, θ::Vector{…})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/xuhUJ/src/hamiltonian.jl:40
 [23] phasepoint(h::AdvancedHMC.Hamiltonian{…}, θ::Vector{…}, r::Vector{…})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/xuhUJ/src/hamiltonian.jl:97
 [24] phasepoint(rng::Random.TaskLocalRNG, θ::Vector{…}, h::AdvancedHMC.Hamiltonian{…})
    @ AdvancedHMC ~/.julia/packages/AdvancedHMC/xuhUJ/src/hamiltonian.jl:179
 [25] initialstep(rng::Random.TaskLocalRNG, model::DynamicPPL.Model{…}, spl::DynamicPPL.Sampler{…}, vi_original::DynamicPPL.TypedVarInfo{…}; initial_params::Nothing, nadapts::Int64, kwargs::@Kwargs{})
    @ Turing.Inference ~/.julia/packages/Turing/IXycm/src/mcmc/hmc.jl:174
 [26] step(rng::Random.TaskLocalRNG, model::DynamicPPL.Model{…}, spl::DynamicPPL.Sampler{…}; initial_params::Nothing, kwargs::@Kwargs{…})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/f9pFn/src/sampler.jl:125
 [27] macro expansion
    @ ~/.julia/packages/AbstractMCMC/7f1oY/src/sample.jl:0 [inlined]
 [28] macro expansion
    @ ~/.julia/packages/AbstractMCMC/7f1oY/src/logging.jl:137 [inlined]
 [29] mcmcsample(rng::Random.TaskLocalRNG, model::DynamicPPL.Model{…}, sampler::DynamicPPL.Sampler{…}, N::Int64; progress::Bool, progressname::String, callback::Nothing, num_warmup::Int64, discard_initial::Int64, thinning::Int64, chain_type::Type, initial_state::Nothing, kwargs::@Kwargs{…})
    @ AbstractMCMC ~/.julia/packages/AbstractMCMC/7f1oY/src/sample.jl:151
 [30] sample(rng::Random.TaskLocalRNG, model::DynamicPPL.Model{…}, sampler::DynamicPPL.Sampler{…}, N::Int64; chain_type::Type, resume_from::Nothing, initial_state::Nothing, progress::Bool, nadapts::Int64, discard_adapt::Bool, discard_initial::Int64, kwargs::@Kwargs{})
    @ Turing.Inference ~/.julia/packages/Turing/IXycm/src/mcmc/hmc.jl:113
 [31] sample
    @ ~/.julia/packages/Turing/IXycm/src/mcmc/hmc.jl:82 [inlined]
 [32] #sample#8
    @ ~/.julia/packages/Turing/IXycm/src/mcmc/Inference.jl:278 [inlined]
 [33] sample
    @ ~/.julia/packages/Turing/IXycm/src/mcmc/Inference.jl:269 [inlined]
 [34] #sample#7
    @ ~/.julia/packages/Turing/IXycm/src/mcmc/Inference.jl:266 [inlined]
 [35] sample(model::DynamicPPL.Model{…}, alg::NUTS{…}, N::Int64)
    @ Turing.Inference ~/.julia/packages/Turing/IXycm/src/mcmc/Inference.jl:263
 [36] top-level scope
    @ ~/Desktop/pkgs/wip/kivu_rp.jl:198

Environment (please complete the following information):

• Output of using Pkg; Pkg.status()

  [31c24e10] Distributions v0.25.120
  [8913a72c] NonlinearSolve v4.10.0
  [fce5fe82] Turing v0.37.0

• Output of versioninfo()

Julia Version 1.10.9
Commit 5595d20a287 (2025-03-10 12:51 UTC)
Build Info:
  Official https://julialang.org/ release
Platform Info:
  OS: macOS (arm64-apple-darwin24.0.0)
  CPU: 8 × Apple M1 Pro
  WORD_SIZE: 64
  LIBM: libopenlibm
  LLVM: libLLVM-15.0.7 (ORCJIT, apple-m1)
Threads: 1 default, 0 interactive, 1 GC (on 6 virtual cores)
Environment:
  JULIA_EDITOR = code
  JULIA_NUM_THREADS = 

[ChrisRackauckas]

You never want to differentiate the solver. This is somewhat difficult to do right in general for ForwardDiff. For arrays at least, https://github.com/SciML/NonlinearSolve.jl/blob/master/lib/BracketingNonlinearSolve/ext/BracketingNonlinearSolveForwardDiffExt.jl does the implicit differentiation rule.

[ayushinav]

I see, though I don't think it's differentiating through the solver, but the coefficients of the equation. It's pretty much like Turing.jl docs on diff eqns here: https://turinglang.org/docs/tutorials/bayesian-differential-equations/ where things work out fine.

[ChrisRackauckas]

Differentiating the coefficients of the equation is differentiating the solver, and that is what the Turing docs are showing. If you change it to an array it will do that just fine.

[ayushinav]

Great, that works out. Thanks!