Introduction of model_typed and model_warntype in DebugUtils

[torfjelde]

I wanted to do some checking of type-stability for a model. Unfortunately, this is quite annoying to do in a way where you can inspect the inferred types of the model's evaluator function.

Hence I was thinking it might be useful to have a few simple methods which just does the set up + calls @code_warntype on model.f (and similarly for @code_typed). This PR therefore introduces model_warntype and model_typed in the DebugUtils module which does exactly this.

I think this is particularly useful for users who wants to check type-stability of their model, which IMO should be able to do so without extensive knowledge of DynamicPPL internals.

The downside: have to add InteractiveUtils.jl as a direct dep (though it's part of base, so should be fine).

Thoughts?

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 84.35%. Comparing base (d6e2147) to head (7f2863b).
Report is 1 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master     #708      +/-   ##
==========================================
+ Coverage   81.90%   84.35%   +2.44%     
==========================================
  Files          30       30              
  Lines        4200     4211      +11     
==========================================
+ Hits         3440     3552     +112     
+ Misses        760      659     -101     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[coveralls]

Pull Request Test Coverage Report for Build 11779219137

Details

• 11 of 11 (100.0%) changed or added relevant lines in 1 file are covered.
• 21 unchanged lines in 2 files lost coverage.
• Overall coverage increased (+2.5%) to 84.351%

---

Files with Coverage Reduction	New Missed Lines	%
src/model.jl	1	93.68%
src/threadsafe.jl	20	50.86%

Totals
Change from base Build 11731628292:	2.5%
Covered Lines:	3552
Relevant Lines:	4211

---

💛 - Coveralls

[yebai]

I think it is very helpful. Also happy to depend on InteractiveUtils.

@willtebbutt you might want to help review this PR too.

[willtebbutt]

I really like this, and agree that InteractiveUtils is a perfectly fine dep to have.

So, I think that the concept is basically fine. Will this stuff automatically appear in the docs, or do we need to add it in manually?

[willtebbutt]

While I think you should feel free to merge this PR without this, could I ask whether you would consider exposing a function which returns the function + arguments that you pass to @code_warntype? i.e. something like

fargs, kwargs = get_args_for_analysis(model, varinfo, context))

with the same default values for varinfo and context as for model_warntype and model_typed?

I guess the implementation would be something like

function get_args_for_analysis(model, varinfo=default_thing, context=default_thing)
    args, kargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
    return (f, args...), kwargs
end

The reason to want this, is that

1. sometimes Base.code_ircode / Base.code_ircode_by_type is the thing you're after, rather than Base.code_typed (the former returns IRCode while the latter returned CodeInfo), and
2. if I want to inspect the IR generated by Mooncake.jl for a particular Turing.jl model (which I often want to do in order to debug stuff), I need to have access to the signature that I'm building the rule for, and I can get this from the arguments that you pass into @code_warntype.

[torfjelde]

Yep, can do @willtebbutt !

[torfjelde]

Wait; why is that function you defined simpler than just doing DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)?

[torfjelde]

Added optimize kwarg to the methods 👍

[willtebbutt]

Wait; why is that function you defined simpler than just doing DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)?

It's a simpler interface, rather than implementation.

It's about ensuring that we have a single source of truth that you can point people towards (and put in the docs) that says "here is everything you need -- we promise that this is the thing that you should profile for performance", as opposed to having to know that the f field of model is the right function to look at (which I assume is an implementation detail rather than a guarantee that the interface makes), and that it must be combined with the output of make_evaluate_arg_and_kwargs.

edit: having now looked at this a bit more closely, am I correct that you can always do:

fargs, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
fargs[1](fargs[2:end]...; kwargs...)

in order to evaluate the correct thing?

[willtebbutt]

Also, I've just realised that this functionality is not tested, and doesn't appear in the docs. It would be good to add a unit test which just runs model_typed and model_warntype in order to check that they do indeed run.

Also, could you please add these to a debug section in the docs?

[torfjelde]

: having now looked at this a bit more closely, am I correct that you can always do:

Nah, it's

model.f(args...; kwargs...)

The evaluator f isn't returned by make_evaluator_args_and_kwargs. I guess returning the evaluator might be useful, but is it clear to have it as part of the "args" as you mentioned? Seems a bit confusing to me; why not return (f, args, kwargs) instead of (f, args...), kwargs?

[torfjelde]

Added docs and testing @willtebbutt

[willtebbutt]

I guess returning the evaluator might be useful, but is it clear to have it as part of the "args" as you mentioned? Seems a bit confusing to me; why not return (f, args, kwargs) instead of (f, args...), kwargs?

I don't feel especially strongly either way. Personally, I like having it with the arguments because for Mooncake.jl-related stuff, where the distinction between function and argument is not really there (e.g. when using Base.code_ircode_by_type you just create the signature for the call). I can also see that most users probably do think about the function as being a separate thing from the args though, so I'm really very happy either way.

[willtebbutt]

This is ready to go once the patch version has been bumped.

As I said, I don't want my other request to get in the way of getting this released. I'm happy to re-review if you do want to implement that as part of this PR though @torfjelde . I'll open an issue about it if you don't fancy doing it as part of this PR.

[willtebbutt]

LGTM. Thanks for indulging my requests!

[willtebbutt]

Sorry, I was hasty. Just realised that there's a problem with the docstring.

[willtebbutt]

This function doesn't actually return the type of the function in question, it returns the function itself. I don't mind which one we do, but this docstring should be fixed before merging!

[willtebbutt]

Might be good to add a unit test to make sure that this is actually doing whatever the docstring asserts is happening.

[torfjelde]

Oh damn, good catch!

[willtebbutt]

Thanks for fixing -- LGTM.