Enzyme: migrate to easy_rule

[wsmoses]

@penelopeysm

[penelopeysm]

For reference, this is the code that this would fix:

using DifferentialEquations, SciMLSensitivity
using ADTypes
using Turing
import Enzyme: set_runtime_activity, Forward, Reverse, Const
using DynamicPPL.TestUtils.AD: run_ad

function ode(du, u, p, t)
    du .= -p .* u
end
prob = ODEProblem(ode, [10.], (0, 10), [2.0])
data = vec(solve(prob, Tsit5(), saveat=0.1)) .+ (randn(101) .* 0.1)

@model function fit_ode(prob)
    σ ~ LogNormal(2,3)
    p ~ truncated(Normal(0., 0.5), lower=0,)

    _prob = remake(prob, p=[p])
    pred = solve(_prob, Tsit5(), saveat=0.1)
    data ~ MvNormal(vec(pred),σ^2 * I)
end

m = fit_ode(prob) | (data = data, )
# s = sample(m, NUTS(), 1000)
run_ad(m, AutoEnzyme(;
        mode = set_runtime_activity(Forward, true),
        function_annotation = Const,
    ))

[penelopeysm]

@wsmoses I wanted to add a test to cover this, but the code above still errors with this PR.

(ppl) pkg> st
Status `~/ppl/Project.toml`
  [47edcb42] ADTypes v1.18.0
  [76274a88] Bijectors v0.15.11 `bi`
  [0c46a032] DifferentialEquations v7.17.0
  [366bfd00] DynamicPPL v0.38.2
  [7da242da] Enzyme v0.13.90
  [12d8515a] EnzymeTestUtils v0.2.5
  [1ed8b502] SciMLSensitivity v7.90.0
  [fce5fe82] Turing v0.41.0
  [8dfed614] Test v1.11.0

on 1.11.7

[wsmoses]

hm whats the error?

[wsmoses]

I'm seeing:

           ))
[ Info: Running AD on fit_ode with AutoEnzyme(mode=EnzymeCore.ForwardMode{false, EnzymeCore.FFIABI, false, true, false}(), function_annotation=Const)
       params : [1.1402281004704176, 0.4386826925439296]

ERROR: MethodError: no method matching forward(::EnzymeCore.EnzymeRules.FwdConfigWidth{2, true, true, true, false}, ::Const{typeof(FastPower.fastpower)}, ::Type{EnzymeCore.BatchDuplicated{Float64, 2}}, ::EnzymeCore.BatchDuplicated{Float64, 2}, ::Const{Float64})

Closest candidates are:
  forward(::EnzymeCore.EnzymeRules.FwdConfig, ::Const{typeof(hypot)}, ::Any, ::EnzymeCore.Annotation, ::EnzymeCore.Annotation, ::EnzymeCore.Annotation, ::EnzymeCore.Annotation...) where N
   @ Enzyme ~/.julia/packages/Enzyme/nOxey/src/internal_rules.jl:1703
  forward(::EnzymeCore.EnzymeRules.FwdConfig, ::Const{Colon}, ::Type{<:Union{EnzymeCore.BatchDuplicated, EnzymeCore.BatchDuplicatedNoNeed, Const, EnzymeCore.Duplicated, EnzymeCore.DuplicatedNoNeed}}, ::EnzymeCore.Annotation{<:AbstractFloat}, ::EnzymeCore.Annotation{<:AbstractFloat}, ::EnzymeCore.Annotation{<:AbstractFloat})
   @ Enzyme ~/.julia/packages/Enzyme/nOxey/src/internal_rules.jl:1305
  forward(::EnzymeCore.EnzymeRules.FwdConfig, ::Const{typeof(Base.range_start_stop_length)}, ::Any, ::EnzymeCore.Annotation{T}, ::EnzymeCore.Annotation{T}, ::EnzymeCore.Annotation{<:Integer}) where T<:Union{Float16, Float32, Float64}
   @ Enzyme ~/.julia/packages/Enzyme/nOxey/src/internal_rules.jl:1197
  ...

Stacktrace:
  [1] getproperty
    @ ~/.julia/packages/SciMLBase/RPW6n/src/integrator_interface.jl:521 [inlined]
  [2] stepsize_controller!
    @ ~/.julia/packages/OrdinaryDiffEqCore/GMkz9/src/integrators/controllers.jl:139 [inlined]
  [3] stepsize_controller!
    @ ~/.julia/packages/OrdinaryDiffEqCore/GMkz9/src/integrators/controllers.jl:5 [inlined]
  [4] _loopfooter!
    @ ~/.julia/packages/OrdinaryDiffEqCore/GMkz9/src/integrators/integrator_utils.jl:267
  [5] getproperty
    @ ~/.julia/packages/SciMLBase/RPW6n/src/integrator_interface.jl:521 [inlined]
  [6] solve!
    @ ~/.julia/packages/OrdinaryDiffEqCore/GMkz9/src/solve.jl:611
  [7] fwddiffe2julia_solve_call_9133_inner_1wrap
    @ ~/.julia/packages/DiffEqBase/p82Yh/src/solve.jl:0
  [8] macro expansion
    @ ~/.julia/packages/Enzyme/nOxey/src/compiler.jl:5883 [inlined]
  [9] enzyme_call
    @ ~/.julia/packages/Enzyme/nOxey/src/compiler.jl:5417 [inlined]
 [10] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/nOxey/src/compiler.jl:5319 [inlined]
 [11] macro expansion
    @ ~/.julia/packages/Enzyme/nOxey/src/rules/jitrules.jl:324 [inlined]
 [12] runtime_generic_fwd(::Type{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::Val{…}, ::typeof(Core.kwcall), ::Nothing, ::Nothing, ::@NamedTuple{…}, ::@NamedTuple{…}, ::@NamedTuple{…}, ::typeof(DiffEqBase.solve_call), ::Nothing, ::Nothing, ::ODEProblem{…}, ::ODEProblem{…}, ::ODEProblem{…}, ::Tsit5{…}, ::Nothing, ::Nothing)
    @ Enzyme.Compiler ~/.julia/packages/Enzyme/nOxey/src/rules/jitrules.jl:357
 [13] #solve_up#39
    @ ~/.julia/packages/DiffEqBase/p82Yh/src/solve.jl:563
 [14] solve_up
    @ ~/.julia/packages/DiffEqBase/p82Yh/src/solve.jl:540 [inlined]
 [15] #solve#38
    @ ~/.julia/packages/DiffEqBase/p82Yh/src/solve.jl:530
 [16] solve
    @ ~/.julia/packages/DiffEqBase/p82Yh/src/solve.jl:520 [inlined]
 [17] fit_ode
    @ ./REPL[12]:6
 [18] _evaluate!!
    @ ~/.julia/packages/DynamicPPL/f38D9/src/model.jl:974 [inlined]
 [19] evaluate_threadunsafe!!
    @ ~/.julia/packages/DynamicPPL/f38D9/src/model.jl:940 [inlined]
 [20] evaluate!!
    @ ~/.julia/packages/DynamicPPL/f38D9/src/model.jl:925 [inlined]
 [21] logdensity_at
    @ ~/.julia/packages/DynamicPPL/f38D9/src/logdensityfunction.jl:242 [inlined]
 [22] logdensity_at
    @ ~/.julia/packages/DynamicPPL/f38D9/src/logdensityfunction.jl:0 [inlined]
 [23] fwddiffe2julia_logdensity_at_8568_inner_1wrap
    @ ~/.julia/packages/DynamicPPL/f38D9/src/logdensityfunction.jl:0
 [24] macro expansion
    @ ~/.julia/packages/Enzyme/nOxey/src/compiler.jl:5883 [inlined]
 [25] enzyme_call
    @ ~/.julia/packages/Enzyme/nOxey/src/compiler.jl:5417 [inlined]
 [26] ForwardModeThunk
    @ ~/.julia/packages/Enzyme/nOxey/src/compiler.jl:5319 [inlined]
 [27] autodiff
    @ ~/.julia/packages/Enzyme/nOxey/src/Enzyme.jl:673 [inlined]
 [28] macro expansion
    @ ~/.julia/packages/Enzyme/nOxey/src/sugar.jl:729 [inlined]
 [29] #gradient#136
    @ ~/.julia/packages/Enzyme/nOxey/src/sugar.jl:585 [inlined]
 [30] value_and_gradient(::typeof(DynamicPPL.logdensity_at), ::DifferentiationInterfaceEnzymeExt.EnzymeForwardGradientPrep{…}, ::AutoEnzyme{…}, ::Vector{…}, ::DifferentiationInterface.Constant{…}, ::DifferentiationInterface.Constant{…}, ::DifferentiationInterface.Constant{…})
    @ DifferentiationInterfaceEnzymeExt ~/.julia/packages/DifferentiationInterface/L0TGS/ext/DifferentiationInterfaceEnzymeExt/forward_onearg.jl:186
 [31] logdensity_and_gradient(f::LogDensityFunction{DynamicPPL.Model{…}, typeof(DynamicPPL.getlogjoint_internal), DynamicPPL.VarInfo{…}, AutoEnzyme{…}}, x::Vector{Float64})
    @ DynamicPPL ~/.julia/packages/DynamicPPL/f38D9/src/logdensityfunction.jl:293
 [32] run_ad(model::DynamicPPL.Model{…}, adtype::AutoEnzyme{…}; test::DynamicPPL.TestUtils.AD.WithBackend{…}, benchmark::Bool, atol::Float64, rtol::Float64, getlogdensity::Function, rng::Random.TaskLocalRNG, varinfo::DynamicPPL.VarInfo{…}, params::Nothing, verbose::Bool)
    @ DynamicPPL.TestUtils.AD ~/.julia/packages/DynamicPPL/f38D9/src/test_utils/ad.jl:270
 [33] run_ad(model::DynamicPPL.Model{typeof(fit_ode), (:prob,), (), (), Tuple{ODEProblem{…}}, Tuple{}, DynamicPPL.ConditionContext{@NamedTuple{…}, DynamicPPL.DefaultContext}}, adtype::AutoEnzyme{EnzymeCore.ForwardMode{false, EnzymeCore.FFIABI, false, true, false}, Const})
    @ DynamicPPL.TestUtils.AD ~/.julia/packages/DynamicPPL/f38D9/src/test_utils/ad.jl:241
 [34] top-level scope
    @ REPL[14]:1
Some type information was truncated. Use `show(err)` to see complete types.

which is an error in a rule outside of here. I'll also fix the rule in that package, but that shows this PR fixes the error the user saw.

[wsmoses]

the latter should be resolved by SciML/FastPower.jl#29

However, the current code here should be considered a non-functional change and can go in independently.

a test confirming it also works more generally can follow the other PR

[penelopeysm]

Thanks @wsmoses! Since this satisfies the tests on 1.10, I'll go ahead and merge it. The tests on 1.11 are still broken, but we'll chat about that separately

[penelopeysm]

Just triggered a new release, so should be out soon!