Julia中ForwardDiff对含Roots.find_zero的函数求导报错如何解决？

你的报错核心是Roots包的find_zero默认会根据传入的初始点类型初始化内部的状态容器，你给的初始点是0.0（Float64类型），但当ForwardDiff运行时，输入的p是携带导数信息的Dual类型，求根过程中产生的中间值也是Dual类型，无法存入固定为Float64的容器，就触发了类型转换错误。

方案1：替换为解析解（适用于你的示例场景） #

你当前示例中的求根逻辑find_zero(z -> logistic(z) - p, 0.0)本质就是求p的对数几率，直接调用StatsFuns自带的logit函数即可，原生支持ForwardDiff微分：

function test_fun(θ::AbstractVector{T}) where T
    μ,σ,p = θ;
    z_star = logit(p)
    return pdf(Normal(μ,σ),z_star)
end
# 此时调用梯度计算不会报错
ForwardDiff.gradient(test_fun,[0.0,1.0,0.75])

方案2：通用求根场景的自定义微分 #

如果你的实际用例中没有解析解，必须用find_zero，可以利用隐函数定理手动定义求根结果的导数，避开让Roots直接处理Dual类型：
对于隐式方程f(z,p)=0，解z(p)对p的导数满足：
dz/dp = - (∂f/∂p) / (∂f/∂z)
你可以先提取p的浮点值完成求根，再手动构造携带导数的Dual类型返回，示例代码如下：

using ForwardDiff, Roots, StatsFuns, Distributions

function solve_z(p::T) where T<:ForwardDiff.Dual
    # 先提取p的数值部分求根
    p_val = ForwardDiff.value(p)
    f(z) = logistic(z) - p_val
    z_val = find_zero(f, 0.0)
    # 计算隐函数导数
    df_dz = logistic(z_val)*(1 - logistic(z_val))
    df_dp = -1
    dz_dp = - df_dp / df_dz
    # 构造带导数的Dual数返回
    return T(z_val, dz_dp * ForwardDiff.partials(p))
end

# 对非Dual类型的p直接正常求根
solve_z(p::Real) = find_zero(z -> logistic(z) - p, 0.0)

function test_fun(θ::AbstractVector{T}) where T
    μ,σ,p = θ;
    z_star = solve_z(p)
    return pdf(Normal(μ,σ),z_star)
end

# 此时梯度计算正常运行
ForwardDiff.gradient(test_fun,[0.0,1.0,0.75])

方案3：适配初始点类型（简单但稳定性弱） #

如果不想手动写导数，也可以把find_zero的初始点改成和输入p同类型的数值，让Roots的内部容器适配Dual类型：

function test_fun(θ::AbstractVector{T}) where T
    μ,σ,p = θ;
    # 初始点用p类型的0，而非固定Float64的0.0
    z_star = find_zero(z -> logistic(z) - p, zero(p))
    return pdf(Normal(μ,σ),z_star)
end

注意：该方案依赖Roots算法的类型兼容性，部分旧版本Roots可能依然会报错，稳定性不如前两种方案。

内容的提问来源于stack exchange，提问作者A. Harris