Forward-over-reverse HVP of logpdf(Dirichlet, x) throws

[rsenne]

Hey @wsmoses this is in reference to #3234.

I forgot to fully explain (and by extension forgot to include the entire MWE) so the llm couldn't see a later issue

I was trying to compute an HVP and the original bounds error that is now fixed in the newer enzyme threw on the gradient calculation. When i however try the HVP calculation i get a new error.

FULL MWE

using Distributions
using Enzyme: Enzyme

f(x) = logpdf(Dirichlet(ones(3)), x)

# in-place reverse-mode gradient of f
function grad!(dx, x)
    dx .= 0.0
    Enzyme.autodiff(Enzyme.Reverse, f, Enzyme.Active, Enzyme.Duplicated(x, dx))
    return nothing
end

x = [0.2, 0.3, 0.5]
v = [1.0, 0.0, -1.0]

println("value        : ", f(x))

dx = zeros(3)
grad!(dx, x)

dx = zeros(3)   
hv = zeros(3)   
vx = copy(v)    
Enzyme.autodiff(
    Enzyme.Forward,
    grad!,
    Enzyme.Duplicated(dx, hv),
    Enzyme.Duplicated(x, vx),
)

Stacktrace

Stacktrace:
  [1] handle_param(args::Vector{…}, codegen_types::Vector{…}, source_typ::Type, rooted_typ::Union{…}, source_i::Int64, orig_i::Int64, arg_jl_i::Int64, codegen_i::Int64, last_cc::GPUCompiler.ArgumentCC, parmsRemoved::Vector{…}, source_sig::Type)
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/typeutils/jltypes.jl:181
  [2] classify_arguments(source_sig::Type, codegen_ft::LLVM.FunctionType, has_sret::Bool, has_returnroots::Bool, has_swiftself::Bool, parmsRemoved::Vector{…}, mi::Core.MethodInstance, world::UInt64)
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/typeutils/jltypes.jl:307
  [3] set_module_types!(interp::Enzyme.Compiler.Interpreter.EnzymeInterpreter{…}, mod::LLVM.Module, primalf::LLVM.Function, job::GPUCompiler.CompilerJob{…}, edges::Vector{…}, run_enzyme::Bool, mode::Enzyme.API.CDerivativeMode)
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:1156
  [4] compile_unhooked(output::Symbol, job::GPUCompiler.CompilerJob{…})
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:5485
  [5] #compile#114
    @ /projectnb/ssmsvi/rsenne/.julia/packages/GPUCompiler/Avrep/src/driver.jl:67 [inlined]
  [6] compile
    @ /projectnb/ssmsvi/rsenne/.julia/packages/GPUCompiler/Avrep/src/driver.jl:55 [inlined]
  [7] _thunk(job::GPUCompiler.CompilerJob{…}, postopt::Bool)
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:6959
  [8] _thunk
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:6957 [inlined]
  [9] cached_compilation
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:7027 [inlined]
 [10] thunkbase(mi::Core.MethodInstance, World::UInt64, FA::Type{…}, A::Type{…}, TT::Type, Mode::Enzyme.API.CDerivativeMode, width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, edges::Vector{…})
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:7143
 [11] thunk_generator(world::UInt64, source::Union{…}, FA::Type, A::Type, TT::Type, Mode::Enzyme.API.CDerivativeMode, Width::Int64, ModifiedBetween::NTuple{…} where N, ReturnPrimal::Bool, ShadowInit::Bool, ABI::Type, ErrIfFuncWritten::Bool, RuntimeActivity::Bool, StrongZero::Bool, self::Any, fakeworld::Any, fa::Type, a::Type, tt::Type, mode::Type, width::Type, modifiedbetween::Type, returnprimal::Type, shadowinit::Type, abi::Type, erriffuncwritten::Type, runtimeactivity::Type, strongzero::Type)
    @ Enzyme.Compiler /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/compiler.jl:7287
 [12] autodiff
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/Enzyme.jl:673 [inlined]
 [13] autodiff
    @ /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/Enzyme.jl:577 [inlined]
 [14] autodiff(::EnzymeCore.ForwardMode{…}, ::var"#grad!", ::EnzymeCore.Duplicated{…}, ::EnzymeCore.Duplicated{…})
    @ Enzyme /projectnb/ssmsvi/rsenne/.julia/packages/Enzyme/Op2XJ/src/Enzyme.jl:549
 [15] top-level scope
    @ /projectnb/ssmsvi/rsenne/ParallelMCMC.jl/dev/mwe_dirichlet_hvp.jl:54
Some type information was truncated. Use `show(err)` to see complete types.