Skip to content

Commit

Permalink
added InteractiveUtils.jl as a dep + model_typed and
Browse files Browse the repository at this point in the history
`model_warntype` for the checking of the model's evaluator
  • Loading branch information
torfjelde committed Oct 31, 2024
1 parent 18af48a commit 9f392a7
Show file tree
Hide file tree
Showing 2 changed files with 34 additions and 0 deletions.
1 change: 1 addition & 0 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
LogDensityProblemsAD = "996a588d-648d-4e1f-a8f0-a84b347e47b1"
Expand Down
33 changes: 33 additions & 0 deletions src/debug_utils.jl
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ using ..DynamicPPL: broadcast_safe, AbstractContext, childcontext

using Random: Random
using Accessors: Accessors
using InteractiveUtils: InteractiveUtils

using DocStringExtensions
using Distributions
Expand Down Expand Up @@ -678,4 +679,36 @@ function has_static_constraints(
return all_the_same(transforms)
end

"""
model_warntype(model[, varinfo, context])
Check the type stability of the model's evaluator, warning about any potential issues.
This simply calls `@code_warntype` on the model's evaluator, filling in internal arguments where needed.
"""
function model_warntype(model::Model, varinfo::AbstractVarInfo=VarInfo(model), context::AbstractContext=DefaultContext())
args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
return if isempty(kwargs)
InteractiveUtils.@code_warntype model.f(args...)
else
InteractiveUtils.@code_warntype model.f(args...; kwargs...)
end
end

"""
model_typed(model[, varinfo, context])
Return the type inference for the model's evaluator.
This simply calls `@code_typed` on the model's evaluator, filling in internal arguments where needed.
"""
function model_typed(model::Model, varinfo::AbstractVarInfo=VarInfo(model), context::AbstractContext=DefaultContext())
args, kwargs = DynamicPPL.make_evaluate_args_and_kwargs(model, varinfo, context)
return if isempty(kwargs)
InteractiveUtils.@code_typed model.f(args...)
else
InteractiveUtils.@code_typed model.f(args...; kwargs...)
end
end

end

0 comments on commit 9f392a7

Please sign in to comment.