Skip to content

Implement Libtask.might_produce_if_sig_contains#218

Merged
penelopeysm merged 3 commits intomainfrom
py/fix-models
Mar 2, 2026
Merged

Implement Libtask.might_produce_if_sig_contains#218
penelopeysm merged 3 commits intomainfrom
py/fix-models

Conversation

@penelopeysm
Copy link
Member

Closes #217. Please see that issue for explanation.

@penelopeysm
Copy link
Member Author

I don't expect this to cause any issues with Libtask. The main question is whether adding this to Turing.jl

Libtask.might_produce_if_sig_contains(::Type{<:DynamicPPL.Model}) = true

causes any slowdown in pMCMC, which I'll test out now.

@github-actions
Copy link

github-actions bot commented Mar 2, 2026

Libtask.jl documentation for PR #218 is available at:
https://TuringLang.github.io/Libtask.jl/previews/PR218/

@penelopeysm penelopeysm merged commit fa0d6e8 into main Mar 2, 2026
22 checks passed
@penelopeysm penelopeysm deleted the py/fix-models branch March 2, 2026 13:22
penelopeysm added a commit to TuringLang/Turing.jl that referenced this pull request Mar 2, 2026
This needs a Libtask release and version bump, which I'll handle once
JuliaRegistrator does its things.

This essentially implements the plan described in
TuringLang/Libtask.jl#217. A lot of the issues
stemming from Libtask not picking up model evaluators either with
keyword arguments, or in submodels, can be fixed by simply declaring
that **every** method that dispatches on `DynamicPPL.Model` is
produceable. The mechanism for this is implemented in
TuringLang/Libtask.jl#218, and this PR makes use
of that.

**For the end-user, this means that we guarantee correctness where
models either have submodels or where models have keyword arguments. The
user no longer has to mark models with keyword arguments as
`@might_produce`.**

I tested performance, and there is no regression — in fact there is a
small speedup (although that is probably benchmark noise):

## Submodel case

This was the issue #2772 where non-inlined submodels were not correctly
picked up.

#2778 fixed this with a strategy that was similar to that in this PR,
but was slightly more limited (this PR handles both submodels and
keyword arguments together).

```julia
using Turing, StableRNGs, Test

@model function inner(y, x)
    @noinline
    y ~ Normal(x)
end
@model function nested(y)
    x ~ Normal()
    a ~ to_submodel(inner(y, x))
end
m1 = nested(1.0)
@time sample(StableRNG(468), m1, PG(10), 2000; chain_type=Any, progress=false);

# 2.585299 seconds on #2778
# 2.523017 seconds on this PR
```

## Keyword argument case

This was the long-standing issue where models with keyword arguments
were originally not picked up by Libtask, and since v0.42.5, could be,
but relied on the user themselves manually declaring
`Libtask.@might_produce`.

```julia
@model function withkw(; y=0.0)
    x ~ Normal()
    y ~ Normal(x)
end
m1 = withkw(y=10.0);
# Turing.@might_produce(withkw)
@time sample(StableRNG(468), m1, PG(10), 2000; chain_type=Any, progress=false);

# withkw case
# 4.741234 seconds on main (requiring @might_produce)
# 4.441797 seconds on this PR (and not requiring @might_produce)
```
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Mark ALL functions that take DynamicPPL.Model as an argument, anywhere, as produceable?

1 participant