Ideally, we should take a more austere approach similar to https://github.com/TuringLang/Turing.jl/issues/2468 and https://github.com/compintell/Mooncake.jl/pull/477