Skip to content

Handle nested PrefixContext #787

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 6 commits into from
Jan 27, 2025
Merged

Handle nested PrefixContext #787

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
01935af
Prefix varnames appropriately inside check_model_and_trace
penelopeysm Jan 25, 2025
e1472ea
Fix values_as_in_model as well
penelopeysm Jan 25, 2025
a76d55d
Add test for check_model with manual prefix
penelopeysm Jan 26, 2025
96afc18
Add values_as_in_model tests
penelopeysm Jan 26, 2025
bbfd011
Add tests for prefix nesting
penelopeysm Jan 26, 2025
47048c2
Bump Project.toml
penelopeysm Jan 26, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 9 additions & 6 deletions src/contexts.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -261,6 +261,7 @@ end

const PREFIX_SEPARATOR = Symbol(".")

# TODO(penelopeysm): Prefixing arguably occurs the wrong way round here
function PrefixContext{PrefixInner}(
context::PrefixContext{PrefixOuter}
) where {PrefixInner,PrefixOuter}
Expand All @@ -273,13 +274,15 @@ function PrefixContext{PrefixInner}(
end
end

function prefix(::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
if @generated
return :(VarName{$(QuoteNode(Symbol(Prefix, PREFIX_SEPARATOR, Sym)))}(getoptic(vn)))
else
VarName{Symbol(Prefix, PREFIX_SEPARATOR, Sym)}(getoptic(vn))
end
# TODO(penelopeysm): Prefixing arguably occurs the wrong way round here
function prefix(ctx::PrefixContext{Prefix}, vn::VarName{Sym}) where {Prefix,Sym}
return prefix(
childcontext(ctx), VarName{Symbol(Prefix, PREFIX_SEPARATOR, Sym)}(getoptic(vn))
)
end
Comment on lines -276 to 282
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

performance

using DynamicPPL, Distributions, Chairmarks

@model function submodel()
    x ~ Normal(0, 1)
end

@model function parentmodel()
    x1 ~ to_submodel(submodel(), true)
end

m = parentmodel()

@be VarInfo(m)
# Current master
Benchmark: 3827 samples with 2 evaluations
 min    10.667 μs (171 allocs: 8.609 KiB)
 median 11.042 μs (171 allocs: 8.609 KiB)
 mean   12.263 μs (171 allocs: 8.609 KiB, 0.03% gc time)
 max    4.467 ms (171 allocs: 8.609 KiB, 98.57% gc time)

# This PR
julia> @be VarInfo(m)
Benchmark: 3749 samples with 2 evaluations
 min    11.209 μs (171 allocs: 8.609 KiB)
 median 11.729 μs (171 allocs: 8.609 KiB)
 mean   12.650 μs (171 allocs: 8.609 KiB, 0.03% gc time)
 max    3.150 ms (171 allocs: 8.609 KiB, 98.04% gc time)

prefix(ctx::AbstractContext, vn::VarName) = prefix(NodeTrait(ctx), ctx, vn)
prefix(::IsLeaf, ::AbstractContext, vn::VarName) = vn
prefix(::IsParent, ctx::AbstractContext, vn::VarName) = prefix(childcontext(ctx), vn)

"""
prefix(model::Model, x)
Expand Down
23 changes: 12 additions & 11 deletions src/debug_utils.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,50 +239,51 @@
end

function record_varname!(context::DebugContext, varname::VarName, dist)
if haskey(context.varnames_seen, varname)
prefixed_varname = prefix(context, varname)
if haskey(context.varnames_seen, prefixed_varname)

Check warning on line 243 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L242-L243 

Added lines #L242 - L243 were not covered by tests
if context.error_on_failure
error("varname $varname used multiple times in model")
error("varname $prefixed_varname used multiple times in model")

Check warning on line 245 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L245 

Added line #L245 was not covered by tests
else
@warn "varname $varname used multiple times in model"
@warn "varname $prefixed_varname used multiple times in model"

Check warning on line 247 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L247 

Added line #L247 was not covered by tests
end
context.varnames_seen[varname] += 1
context.varnames_seen[prefixed_varname] += 1

Check warning on line 249 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L249 

Added line #L249 was not covered by tests
else
# We need to check:
# 1. Does this `varname` subsume any of the other keys.
# 2. Does any of the other keys subsume `varname`.
vns = collect(keys(context.varnames_seen))
# Is `varname` subsumed by any of the other keys?
idx_parent = findfirst(Base.Fix2(subsumes, varname), vns)
idx_parent = findfirst(Base.Fix2(subsumes, prefixed_varname), vns)

Check warning on line 256 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L256 

Added line #L256 was not covered by tests
if idx_parent !== nothing
varname_parent = vns[idx_parent]
if context.error_on_failure
error(
"varname $(varname_parent) used multiple times in model (subsumes $varname)",
"varname $(varname_parent) used multiple times in model (subsumes $prefixed_varname)",
)
else
@warn "varname $(varname_parent) used multiple times in model (subsumes $varname)"
@warn "varname $(varname_parent) used multiple times in model (subsumes $prefixed_varname)"

Check warning on line 264 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L264 

Added line #L264 was not covered by tests
end
# Update count of parent.
context.varnames_seen[varname_parent] += 1
else
# Does `varname` subsume any of the other keys?
idx_child = findfirst(Base.Fix1(subsumes, varname), vns)
idx_child = findfirst(Base.Fix1(subsumes, prefixed_varname), vns)

Check warning on line 270 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L270 

Added line #L270 was not covered by tests
if idx_child !== nothing
varname_child = vns[idx_child]
if context.error_on_failure
error(
"varname $(varname_child) used multiple times in model (subsumed by $varname)",
"varname $(varname_child) used multiple times in model (subsumed by $prefixed_varname)",
)
else
@warn "varname $(varname_child) used multiple times in model (subsumed by $varname)"
@warn "varname $(varname_child) used multiple times in model (subsumed by $prefixed_varname)"

Check warning on line 278 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L278 

Added line #L278 was not covered by tests
end

# Update count of child.
context.varnames_seen[varname_child] += 1
end
end

context.varnames_seen[varname] = 1
context.varnames_seen[prefixed_varname] = 1

Check warning on line 286 in src/debug_utils.jl

View check run for this annotation

Codecov / codecov/patch

src/debug_utils.jl#L286 

Added line #L286 was not covered by tests
end
end

Expand Down
2 changes: 1 addition & 1 deletion src/values_as_in_model.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,7 @@ is_extracting_values(::IsParent, ::AbstractContext) = false
is_extracting_values(::IsLeaf, ::AbstractContext) = false

function Base.push!(context::ValuesAsInModelContext, vn::VarName, value)
return setindex!(context.values, copy(value), vn)
return setindex!(context.values, copy(value), prefix(context, vn))
end

function broadcast_push!(context::ValuesAsInModelContext, vns, values)
Expand Down
Loading