Fix inference failure in callbacks

[charleskawczynski]

This PR fixes the broken inference test for the callback loop. Closes #316 and #315.

[Sbozzolo]

How does this affect compilation times?

[charleskawczynski]

How does this affect compilation times?

Better inferred code correlates with shorter compilation time, but heuristics are involved, so we’d need to test it out to know for sure.

My guess is that this shouldn’t have a notable impact.

[Sbozzolo]

Better inferred code correlates with shorter compilation time, but heuristics are involved, so we’d need to test it out to know for sure.

Maybe I should have said inference + compilation. Better inferred code can lead to longer inference + compilation time (it's much easier to propagate Any then solving the inference problems). I've also seen places where unrolling led to subtantial increase in latency.

My guess is that this shouldn’t have a notable impact.

I also think so because we don't have many callbacks, but it would be great if you could double-check this before merging.

[charleskawczynski]

One reason to continue with this is that, sometimes, JET will show inference failures up to a point, and bail on showing more. This inference failure seems to be the last one seen in a handful of jobs in CliMA/ClimaAtmos.jl#3290. It very well may be that this is the last one, however, it’s possible that it’s the last failure before a bunch more, which is what I’d like to know.

Ultimately, fixing inference failures is a sort of due diligence in that latency is often attributed to poorly inferred code. I will say though, CliMA/ClimaAtmos.jl#3290 is increasingly making me think that that JuliaLang/julia#55807 is the real issue.

[charleskawczynski]

I also think so because we don't have many callbacks, but it would be great if you could double-check this before merging.

sure, I’ll run clima atmos with this branch tomorrow

[charleskawczynski]

Build here: https://buildkite.com/clima/climaatmos-ci/builds/20799.

[charleskawczynski]

(cc @Sbozzolo)

[Sbozzolo]

Thank you!