Skip to content

Commit 7a8e7e3

Browse files
committed
Remove SamplingContext from JETExt to avoid triggering Sampling->Init pathway
1 parent dfbf726 commit 7a8e7e3

File tree

1 file changed

+5
-10
lines changed

1 file changed

+5
-10
lines changed

ext/DynamicPPLJETExt.jl

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -21,22 +21,17 @@ end
2121
function DynamicPPL.Experimental._determine_varinfo_jet(
2222
model::DynamicPPL.Model; only_ddpl::Bool=true
2323
)
24-
# Use SamplingContext to test type stability.
25-
sampling_model = DynamicPPL.contextualize(
26-
model, DynamicPPL.SamplingContext(model.context)
27-
)
28-
2924
# First we try with the typed varinfo.
30-
varinfo = DynamicPPL.typed_varinfo(sampling_model)
25+
varinfo = DynamicPPL.typed_varinfo(model)
3126

32-
# Let's make sure that both evaluation and sampling doesn't result in type errors.
27+
# Let's make sure that evaluation doesn't result in type errors.
3328
issuccess, result = DynamicPPL.Experimental.is_suitable_varinfo(
34-
sampling_model, varinfo; only_ddpl
29+
model, varinfo; only_ddpl
3530
)
3631

3732
if !issuccess
3833
# Useful information for debugging.
39-
@debug "Evaluaton with typed varinfo failed with the following issues:"
34+
@debug "Evaluation with typed varinfo failed with the following issues:"
4035
@debug result
4136
end
4237

@@ -46,7 +41,7 @@ function DynamicPPL.Experimental._determine_varinfo_jet(
4641
else
4742
# Warn the user that we can't use the type stable one.
4843
@warn "Model seems incompatible with typed varinfo. Falling back to untyped varinfo."
49-
DynamicPPL.untyped_varinfo(sampling_model)
44+
DynamicPPL.untyped_varinfo(model)
5045
end
5146
end
5247

0 commit comments

Comments
 (0)