Help to use AD with Optimization.jl

[gdalle]

@Maucejo

With ForwardDiff.jl (Optimization.AutoForwardDiff()), I get the standard error

ERROR: MethodError: no method matching Float64(::ForwardDiff.Dual{ForwardDiff.Tag{typeof(optimQ!), Float64}, Float64, 2})
The type `Float64` exists, but no method is defined for this combination of argument types when trying to construct it.

As you can see in the code, I tried to make the code generic (it is type stable), but it doesn’t work and I don’t know why

That is because you pre-allocated storage with Float64 element type, which during optimization will be populated with Dual numbers. The solution for that is to handle your storage with PreallocationTools.jl, until Optimization.jl can make better use of the Cache objects provided by DifferentiationInterface.jl.

@Maucejo

With Zygote.jl (Optimization.AutoZygote()), I get this error:

 Mutating arrays is not supported -- called copyto!(Matrix{Float64}, ...)
 This error occurs when you ask Zygote to differentiate operations that change
 the elements of arrays in place (e.g. setting values with x .= ...)

I understand this error, because I use preallocations and broadcasting (but how to avoid this).

To avoid this error, you would need to define a custom rule with ChainRulesCore.jl for your high-level function as a whole. That is what I did for HiddenMarkovModels.jl, where it turns out that the forward-backward algorithm gives you all the necessary information for derivatives. I assume similar results exist for the Kalman filter.
Alternately, Zygote.Buffer might be able to handle the simplest in-place operations.

@Maucejo

With Enzyme.jl (Optimization.AutoEnzyme()), I get this warning:

TODO forward zero-set of memory copy used memset rather than runtime type

and finally this error:

ERROR: Enzyme compilation failed due to illegal type analysis.
This usually indicates the use of a Union type, which is not fully supported with Enzyme.API.strictAliasing set to true [the default].
Ideally, remove the union (which will also make your code faster), or try setting Enzyme.API.strictAliasing!(false) before any autodiff call.
To toggle more information for debugging (needed for bug reports), set Enzyme.Compiler.VERBOSE_ERRORS[] = true (default false)

The only Union I have is the signature of the for loop.

Have you tried activating the recommended Enzyme.jl options? What is the stack trace?

---

Also, have you tried Mooncake.jl with AutoMooncake(; config=nothing)?