Type stable code with autodiff

General Usage
1

Dear Julia users

I am writing a quite important julia project where performances are very important, and, to enforce type stability, almost all my functions look like this:

function get_kinetic_energy(scha::SCHA{T}, ensemble::Ensemble{T}, nltransf::NLTransform{T}; 
        debug::Bool=false)::T where {T}

I thought it was a best practice to tell beforehand to the compiler that all the types of the things inside are the same, so no type conversion is required when there are operations between the variables.

However, now I am experimenting with automatic differentiation, but I realized that the type of variables I want to differentiate is the ForwardDiff.Dual, and this prevents the multiple dispatchers to correctly resolving the types (as some types are T (Float64), others are Dual).
Indeed, if I remove all the type constraints, the autodiff works perfectly, but I fear it would lose performance by not allowing the compiler to fix potential type instabilities.

What are the best practices for writing code that is both automatically differentiable and type-stable?

2

This has nothing to do with type stability. Dispatches auto-specialize on types.

3

Type annotations will not cause code to become type stable, they only restrict what input types are allowed. If your code is type stable with annotations, it will also be type stable without them, for the same inputs, at least.

You can use other tools to check for type stability. Traditionally, people have used @code_warntype for this, but now there are some newer ones, like Cthulu.jl, I believe.

3 Likes