diff --git a/HISTORY.md b/HISTORY.md index f853ff1be..2293fab24 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -142,6 +142,14 @@ Handwritten evaluators must prepare argument bindings: direct argument use → r Debug introspection may return a body callable distinct from `model.f`: assuming `gen_evaluator_call_with_types(m)[1] === m.f` → use the returned callable and argument types. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501). +`PrefixContext` is removed: context-based prefixing → `prefix(model, vn; template)`. See [#1502](https://github.com/TuringLang/DynamicPPL.jl/pull/1502), [#1221](https://github.com/TuringLang/DynamicPPL.jl/issues/1221). + +`extract_prefixes` is removed: `extract_prefixes(model.context)` → `DynamicPPL.getprefix(model)`, returning the combined prefix or `nothing`. See [#1502](https://github.com/TuringLang/DynamicPPL.jl/pull/1502). + +Context-based address prefixing is removed: `prefix(ctx, vn)` → `AbstractPPL.prefix(vn, DynamicPPL.getprefix(model))` for prefixed models. See [#1502](https://github.com/TuringLang/DynamicPPL.jl/pull/1502). + +Bindings outside a model’s prefix now throw: `condition(prefix(m, :p); y=v)` → `condition(prefix(m, :p), @varname(p.y) => v)`. See [#1502](https://github.com/TuringLang/DynamicPPL.jl/pull/1502). + `@vnt` is no longer exported: `@vnt` → `DynamicPPL.@vnt` or explicitly import it. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501). `CondFixContext` is removed: context-stored observations/fixed bindings → `condition(model, values)`/`fix(model, values)`. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501). @@ -156,7 +164,7 @@ Context binding accessors are removed: `conditioned(context)`/`fixed(context)` Context observation hooks are removed: `tilde_observe!!(::AbstractContext, ...)` overloads → `accumulate_observe!!` implementations. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501). -Observation calls change: `tilde_observe!!(ctx, dist, value, vn, template, vi)` → `tilde_observe!!(prefix, prefix_template, dist, value, vn, template, vi)`. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501). +Observation calls take the model's prefix instead of a context: `tilde_observe!!(ctx, dist, value, vn, template, vi)` → `tilde_observe!!(prefix, dist, value, vn, template, vi)`. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501), [#1502](https://github.com/TuringLang/DynamicPPL.jl/pull/1502). Submodel latent calls require the parent model: `tilde_assume!!(ctx, submodel, vn, template, vi)` → `tilde_assume!!(parent, ctx, submodel, vn, template, vi)`. See [#1501](https://github.com/TuringLang/DynamicPPL.jl/pull/1501). diff --git a/docs/src/api.md b/docs/src/api.md index dfea12e17..c1ded5d40 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -498,7 +498,7 @@ InitContext Customise latent value selection with an [initialisation strategy](init.md) supplied to `InitContext`. `tilde_assume!!` dispatches on that context to initialise, transform, and accumulate a latent value. -Every observation calls `tilde_observe!!(prefix, prefix_template, right, left, vn, template, vi)`, which applies the prefix metadata and calls `accumulate_observe!!` without dispatching on a context. +Every observation calls `tilde_observe!!(prefix, right, left, vn, template, vi)`, which applies the prefix metadata and calls `accumulate_observe!!` without dispatching on a context. ```@docs tilde_assume!! @@ -506,7 +506,7 @@ tilde_observe!! ``` **Parent contexts**: These essentially act as 'modifiers' for leaf contexts. -`PrefixContext` supplies address metadata. Conditioned and fixed values are stored on the model. +Prefixes, conditioned values, and fixed values are stored on the model. To implement a parent context, you have to subtype `DynamicPPL.AbstractParentContext`, and implement the `childcontext` and `setchildcontext` methods. If needed, you can also implement `tilde_assume!!` for your context. @@ -526,12 +526,6 @@ leafcontext setleafcontext ``` -Sometimes it is necessary to handle all `PrefixContext`s in a context stack at one go: - -```@docs -extract_prefixes -``` - ### VarInfo initialisation The function `init!!` is used to initialise, or overwrite, values in a VarInfo. diff --git a/docs/src/conditionfix.md b/docs/src/conditionfix.md index 0409829f8..aa09220ca 100644 --- a/docs/src/conditionfix.md +++ b/docs/src/conditionfix.md @@ -435,6 +435,9 @@ and complete local LHS variables retain integer indices. A **submodel namespace* child LHS variables through addresses such as `a.x`, or through unchanged names with `auto_prefix=false` unless manually prefixed. +Bindings on prefixed models use absolute addresses: for `prefix(m, @varname(p))`, +use `@varname(p.y) => 0.0`; binding `y=0.0` throws `ArgumentError`. + For example, if `a` is an argument, `a ~ to_submodel(child())`, `a[1] ~ to_submodel(child())` and `a.x ~ to_submodel(child())` are rejected by the **Submodel LHS and prefixes** rule. Use a new local name and condition the child's LHS diff --git a/docs/src/migration.md b/docs/src/migration.md index a78b2cb81..02dfea35f 100644 --- a/docs/src/migration.md +++ b/docs/src/migration.md @@ -12,6 +12,16 @@ or `VectorValueAccumulator` when those outputs are needed, for example `VarInfo(VectorValueAccumulator(), DynamicPPL.default_accumulators()...)`; `VarInfo(model)` remains a convenience constructor that records vectorised values and log densities. +`PrefixContext` and `extract_prefixes` are removed. Use `prefix(model, vn; template)` +to set a prefix and `DynamicPPL.getprefix(model)` to read the combined prefix +(`nothing` when absent). The `prefix` field stores internal metadata for LHS variable +addresses and nested submodel namespace storage templates. Replace `prefix(ctx, vn)` with +`AbstractPPL.prefix(vn, DynamicPPL.getprefix(model))` for a prefixed model. + +Bindings outside a model's prefix now throw. Replace +`condition(prefix(m, @varname(p)); y=v)` with +`condition(prefix(m, @varname(p)), @varname(p.y) => v)`. + To reuse previous values, extract them explicitly before evaluating: ```julia @@ -346,7 +356,8 @@ before replaying them. Replace context overloads of `tilde_observe!!` with `accumulate_observe!!` implementations. For direct observation calls, replace `tilde_observe!!(ctx, dist, value, vn, template, vi)` -with the current signature `tilde_observe!!(prefix, prefix_template, dist, value, vn, template, vi)`. +with `tilde_observe!!(prefix, dist, value, vn, template, vi)`, where `prefix` is the model's +prefix value (`nothing`, a `VarName`, or a `PrefixTemplate`). For submodel latent calls, replace `tilde_assume!!(ctx, submodel, vn, template, vi)` with `tilde_assume!!(parent, ctx, submodel, vn, template, vi)`. Replace `store_coloneq_value!!(ctx, vn, value, template, vi)` with diff --git a/docs/src/tilde.md b/docs/src/tilde.md index 4bba315c6..672eebd95 100644 --- a/docs/src/tilde.md +++ b/docs/src/tilde.md @@ -33,7 +33,7 @@ begin # Conceptual branch: the compiler knows whether x is an argument. conditioned_x = is_argument(x) ? x : get_conditioned_value(__model__, vn) raw_x, __varinfo__ = tilde_observe!!( - __prefix__, __prefix_template__, dist, conditioned_x, vn, template, __varinfo__ + __prefix__, dist, conditioned_x, vn, template, __varinfo__ ) else @@ -46,8 +46,9 @@ begin end ``` -Here, `__prefix__` and `__prefix_template__` are the model's prefix and its storage -template, passed directly to the exported `tilde_observe!!` entry point. +Here, `__prefix__` is the model's single prefix value (`__model__.prefix`): `nothing`, +a `VarName`, or a `PrefixTemplate` containing storage metadata. It is passed directly +to the exported `tilde_observe!!` entry point. We won't go into detail about every part of this code; by far the most interesting part is the call to `tilde_assume!!`. Every latent tilde-statement `vn ~ dist` is transformed into one such call. diff --git a/docs/src/to_distribution_and_to_submodel.md b/docs/src/to_distribution_and_to_submodel.md index 174addcef..3a2e322d0 100644 --- a/docs/src/to_distribution_and_to_submodel.md +++ b/docs/src/to_distribution_and_to_submodel.md @@ -60,4 +60,5 @@ initializer in unconstrained space. to_distribution to_submodel DynamicPPL.prefix +DynamicPPL.getprefix ``` diff --git a/ext/DynamicPPLInputProvenanceExt.jl b/ext/DynamicPPLInputProvenanceExt.jl index 5b8426d09..e8f438ca9 100644 --- a/ext/DynamicPPLInputProvenanceExt.jl +++ b/ext/DynamicPPLInputProvenanceExt.jl @@ -195,8 +195,9 @@ function check_input_provenance(rng, model, params) model.f, args, defaults, - model.context, - values; + model.prefix, + values, + model.context; args_on_lhs=DynamicPPL._args_on_lhs(model), ) vi = DynamicPPL.VarInfo((InputProvenanceAccumulator(),)) diff --git a/src/DynamicPPL.jl b/src/DynamicPPL.jl index ad39072f9..257301a40 100644 --- a/src/DynamicPPL.jl +++ b/src/DynamicPPL.jl @@ -150,7 +150,6 @@ export AbstractVarInfo, setchildcontext, leafcontext, setleafcontext, - extract_prefixes, # Tilde pipeline tilde_assume!!, tilde_observe!!, @@ -226,6 +225,8 @@ export AbstractVarInfo, generated_quantities, typed_identity +@compat public getprefix + # Reexport using Distributions: loglikelihood export loglikelihood @@ -264,7 +265,6 @@ include("transformed_values.jl") include("contexts.jl") include("contexts/default.jl") include("contexts/init.jl") -include("contexts/prefix.jl") include("model.jl") include("conditionfix.jl") @compat public Recursive diff --git a/src/compiler.jl b/src/compiler.jl index 7b4a67d71..664af6b75 100644 --- a/src/compiler.jl +++ b/src/compiler.jl @@ -407,8 +407,7 @@ function generate_tilde_literal(left, right) @gensym value return quote $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)( - $(DynamicPPL._model_prefix)(__model__), - $(DynamicPPL._model_prefix_template)(__model__), + __model__.prefix, $(DynamicPPL.check_tilde_rhs)($right), $left, nothing, @@ -530,7 +529,7 @@ function generate_tilde(left, right; is_argument=false) if $role isa $(DynamicPPL.Fix) $value = $(DynamicPPL._check_tilde_value)( $fixed_data, - $(DynamicPPL.maybe_prefix)($vn, $(DynamicPPL._model_prefix)(__model__)), + $(DynamicPPL.maybe_prefix)($vn, $(DynamicPPL.getprefix)(__model__)), $role, ) $(assign_or_set!!(left, value, vn, is_argument ? _set_fixed_lhs : _set_lhs)) @@ -558,8 +557,7 @@ function generate_tilde(left, right; is_argument=false) ) $value, __varinfo__ = $(DynamicPPL.tilde_observe!!)( - $(DynamicPPL._model_prefix)(__model__), - $(DynamicPPL._model_prefix_template)(__model__), + __model__.prefix, $(DynamicPPL.check_tilde_rhs)($dist), $supplied_val, $vn, @@ -603,7 +601,7 @@ function generate_input_provenance_check(left::Union{Expr,Symbol}, vn) __varinfo__ = $(DynamicPPL.check_input_provenance!!)( __varinfo__, $value, - $(DynamicPPL.maybe_prefix)($vn, $(DynamicPPL._model_prefix)(__model__)), + $(DynamicPPL.maybe_prefix)($vn, $(DynamicPPL.getprefix)(__model__)), ) end end @@ -768,7 +766,7 @@ function build_output( # See the docstrings of `replace_returns` for more info. evaluatordef[:body] = MacroTools.@q begin $(linenumbernode) - __context__ = first($(DynamicPPL.extract_prefixes)(__model__.context)) + __context__ = __model__.context $(replace_returns(add_return_to_last_statment(modeldef[:body]))) end @@ -938,7 +936,7 @@ end function prepare_model_argument(model::Model, vn::VarName, value) binding = _get_model_binding(model, vn) - return _model_argument_value(binding, value, maybe_prefix(vn, _model_prefix(model))) + return _model_argument_value(binding, value, maybe_prefix(vn, getprefix(model))) end function prepare_model_argument(binding, value) return _model_argument_value(binding, value) diff --git a/src/conditionfix.jl b/src/conditionfix.jl index 235bb6761..f06819a04 100644 --- a/src/conditionfix.jl +++ b/src/conditionfix.jl @@ -420,7 +420,7 @@ end end function _get_model_role(model, vn, template=NoTemplate()) root = _model_value_varname( - model.values, VarName{AbstractPPL.getsym(vn)}(), _model_prefix(model) + model.values, VarName{AbstractPPL.getsym(vn)}(), getprefix(model) ) if template isa NamedTuple binding = _model_argument_binding( @@ -436,7 +436,7 @@ function _get_model_role(model, vn, template=NoTemplate()) end end end - vn = _model_value_varname(model.values, vn, _model_prefix(model)) + vn = _model_value_varname(model.values, vn, getprefix(model)) return _model_role_at( _model_values(model.values), AbstractPPL.varname_to_optic(vn), @@ -444,7 +444,7 @@ function _get_model_role(model, vn, template=NoTemplate()) ) end function _get_model_binding(model, vn) - vn = _model_value_varname(model.values, vn, _model_prefix(model)) + vn = _model_value_varname(model.values, vn, getprefix(model)) return _model_argument_binding( _model_values(model.values), AbstractPPL.varname_to_optic(vn) ) @@ -455,7 +455,7 @@ function _get_argument_role(model, vn, argument) return binding isa ModelValue ? _model_role(binding, vn) : _get_model_role(model, vn) end function _get_model_data(model, vn) - vn = _model_value_varname(model.values, vn, _model_prefix(model)) + vn = _model_value_varname(model.values, vn, getprefix(model)) binding = _model_argument_binding( _model_values(model.values), AbstractPPL.varname_to_optic(vn) ) @@ -473,7 +473,7 @@ function _get_model_data(model, vn, argument, local_value) AbstractPPL.getoptic(vn), vn, _fixed_owners(model.values), - _model_value_varname(model.values, argument, _model_prefix(model)), + _model_value_varname(model.values, argument, getprefix(model)), ) return _get_model_data(model, vn) end @@ -1406,7 +1406,7 @@ function _check_latent_arguments(model) AbstractPPL.getsym(vn) in lhs || return nothing binding = _get_model_binding(model, vn) storage = binding === nothing ? value : _argument_storage(binding, value) - _check_latent_storage(storage, maybe_prefix(vn, _model_prefix(model))) + _check_latent_storage(storage, maybe_prefix(vn, getprefix(model))) end return model end @@ -1712,7 +1712,7 @@ end @generated function _argument_may_alias(::Type{T}) where {T} return _argument_alias_type(T, Set{Any}()) end -function _check_shared_latent_storage(model::Model, prefix=_model_prefix(model)) +function _check_shared_latent_storage(model::Model, prefix=getprefix(model)) lhs = _args_on_lhs(model) latent, repeated = Base.IdSet{Any}(), Any[] roots, bound = Tuple{Any,Any,Any,Union{Bool,Nothing}}[], Any[] @@ -2135,8 +2135,8 @@ function _materialize_argument_values(model::Model) model.values isa UnprefixedArgumentValues || return model values = _prefix_values( _model_values(model.values), - _model_prefix(model), - _apply_prefix_template(_model_prefix_template(model), NoTemplate()), + getprefix(model), + _apply_prefix_template(model.prefix, NoTemplate()), ) return _reconstruct_model(model; values) end @@ -2155,7 +2155,7 @@ function _check_binding_addresses(model, values) metadata = _binding_metadata(model) names = _lhs_names(metadata) names === nothing && return nothing - prefix = model.values isa LocalModelValues ? nothing : _model_prefix(model) + prefix = model.values isa LocalModelValues ? nothing : getprefix(model) if prefix !== nothing for name in keys(values.data) name === AbstractPPL.getsym(prefix) || throw( @@ -2187,7 +2187,7 @@ end quote name = $(QuoteNode(name)) stored_name = $(QuoteNode(stored_name)) - vn = _model_value_varname(model.values, VarName{name}(), _model_prefix(model)) + vn = _model_value_varname(model.values, VarName{name}(), getprefix(model)) binding = _model_argument_binding(values, AbstractPPL.varname_to_optic(vn)) if name in _args_on_lhs(model) if binding isa Union{VarNamedTuple,VarNamedTuples.PartialArray} @@ -2729,7 +2729,7 @@ function _prepare_local_binding_children( end end function _prepare_local_binding_types(::Type{R}, model, values) where {R} - prefix = _model_prefix(model) + prefix = getprefix(model) local_values = if model.values isa LocalModelValues || prefix === nothing values else @@ -2755,9 +2755,7 @@ function _prepare_local_binding_types(::Type{R}, model, values) where {R} local_values else _prefix_values( - local_values, - prefix, - _apply_prefix_template(_model_prefix_template(model), NoTemplate()), + local_values, prefix, _apply_prefix_template(model.prefix, NoTemplate()) ) end end @@ -2803,7 +2801,7 @@ function _surviving_binding_templates(previous, values, ::Type{R}, replaced=()) end end function _submodel_binding_templates(model, prefix) - prefix = _model_value_varname(model.values, prefix, _model_prefix(model)) + prefix = _model_value_varname(model.values, prefix, getprefix(model)) return mapreduce((a, b) -> (a..., b...), _binding_templates(model.values); init=()) do t prefix === nothing && return (t,) subsumes(prefix, t.name) || return () @@ -2917,7 +2915,7 @@ function _bind_inputs(::Type{R}, model::Model, inputs::Tuple) where {R} _check_shared_latent_storage(model) isempty(deferred) && return model templates = _binding_template_entries( - R, deferred, model.values isa LocalModelValues ? nothing : _model_prefix(model) + R, deferred, model.values isa LocalModelValues ? nothing : getprefix(model) ) bindings = _with_binding_templates( model.values, _merge_binding_templates(_binding_templates(model.values), templates) @@ -2931,7 +2929,7 @@ function _bound_binding_schema(model, schema::NamedTuple, values, prefix=nothing if storage isa NamedTuple _bound_binding_schema(model, storage, values, address) else - root = _model_value_varname(model.values, address, _model_prefix(model)) + root = _model_value_varname(model.values, address, getprefix(model)) any(v -> _input_binds_root(v, root, model), values) ? storage : nothing end end @@ -2974,7 +2972,7 @@ function _binding_template_names(template::NamedTuple, prefix=nothing) end end function _record_binding_template(model::Model, template) - prefix = _model_prefix(model) + prefix = getprefix(model) names = map(_binding_template_names(template)) do name local_name = prefix === nothing || (name != prefix && subsumes(prefix, name)) address = @@ -2986,14 +2984,15 @@ function _record_binding_template(model::Model, template) model.f, model.args, model.defaults, - model.context, - model.values; + model.prefix, + model.values, + model.context; args_on_lhs=metadata, ) end function _local_binding_schema(model, schema, values) - prefix = _model_prefix(model) + prefix = getprefix(model) prefix === nothing && return schema local_schema = _schema_namespace(schema, AbstractPPL.varname_to_optic(prefix)) for name in keys(schema) @@ -3021,15 +3020,15 @@ function _schema_namespace(value::VarNamedTuple, optic::AbstractPPL.Property) end function _binding_display_name(model, vn) - return model.values isa LocalModelValues ? maybe_prefix(vn, _model_prefix(model)) : vn + return model.values isa LocalModelValues ? maybe_prefix(vn, getprefix(model)) : vn end function _schema_binding_address(model, vn) vn = _expand_cartesian(vn) - prefix = _model_prefix(model) + prefix = getprefix(model) if !(model.values isa LocalModelValues) && prefix !== nothing && AbstractPPL.getsym(prefix) === AbstractPPL.getsym(vn) - template = _apply_prefix_template(_model_prefix_template(model), NoTemplate()) + template = _apply_prefix_template(model.prefix, NoTemplate()) return _concretize_prefix( vn, template; depth=optic_skip_length(AbstractPPL.getoptic(prefix)) ) @@ -3138,7 +3137,7 @@ function _prepare_schema_input( layer = _model_values(layer_model.values) templates = VarNamedTuple() for (name, storage) in pairs(schema) - root = _model_value_varname(model.values, VarName{name}(), _model_prefix(model)) + root = _model_value_varname(model.values, VarName{name}(), getprefix(model)) _input_binds_root(input, root, model) || continue _check_deferred_schema(storage, root, input, model) previous = _model_argument_binding(layer, AbstractPPL.varname_to_optic(root)) @@ -3247,11 +3246,11 @@ function _bind_model(::Type{R}, model::Model, values; preparation_model=model) w end, ), ) - if !(model.values isa LocalModelValues) && _model_prefix(model) !== nothing + if !(model.values isa LocalModelValues) && getprefix(model) !== nothing owners = _prefix_values( owners, - _model_prefix(model), - _apply_prefix_template(_model_prefix_template(model), NoTemplate()), + getprefix(model), + _apply_prefix_template(model.prefix, NoTemplate()), ) end owners = VarNamedTuple(merge(owners.data, fixed_values.data)) @@ -3365,7 +3364,7 @@ end arguments = map(fields) do stored_name name = unsplat_symbol(stored_name) quote - root = _model_value_varname(model.values, $(VarName{name}()), _model_prefix(model)) + root = _model_value_varname(model.values, $(VarName{name}()), getprefix(model)) if subsumes(root, address) argument = merge(model.args, model.defaults)[$(QuoteNode(stored_name))] previous = _model_argument_binding( @@ -3425,7 +3424,7 @@ function _check_local_property_index(model, vn; operation="bind") isempty(fields) && return nothing template = NamedTuple{fields}(map(_ -> NoTemplate(), fields)) # Keep the namespace, replacing only the local root's index path. - root = _model_value_varname(model.values, VarName{name}(), _model_prefix(model)) + root = _model_value_varname(model.values, VarName{name}(), getprefix(model)) return _check_namedtuple_index( template, optic, AbstractPPL.varname_to_optic(root); operation ) @@ -3446,7 +3445,7 @@ end for stored_name in keys(model.defaults) is_splat_symbol(stored_name) || continue argument = unsplat_symbol(stored_name) - root = _model_value_varname(model.values, VarName{argument}(), _model_prefix(model)) + root = _model_value_varname(model.values, VarName{argument}(), getprefix(model)) if root != vn && subsumes(root, vn) if operation == "remove" optic = AbstractPPL._unprefix_optic( @@ -3488,13 +3487,13 @@ end return vn, template end function _binding_template( - model, templates::VarNamedTuple, vn::VarName; prefix=_model_prefix(model) + model, templates::VarNamedTuple, vn::VarName; prefix=getprefix(model) ) template = get(templates.data, AbstractPPL.getsym(vn), NoTemplate()) if template isa NoTemplate && prefix !== nothing && AbstractPPL.getsym(vn) === AbstractPPL.getsym(prefix) - return _apply_prefix_template(_model_prefix_template(model), NoTemplate()) + return _apply_prefix_template(model.prefix, NoTemplate()) end return template end @@ -4264,7 +4263,7 @@ function _remove_marked(::Type{R}, values, r::ModelRemoval) where {R} end function _local_removal_name(model, vn; operation="remove") - prefix = model.values isa LocalModelValues ? nothing : _model_prefix(model) + prefix = model.values isa LocalModelValues ? nothing : getprefix(model) if vn === nothing || prefix === nothing return vn elseif subsumes(prefix, vn) @@ -4378,7 +4377,7 @@ function _select_removal(r, prefix) return (ModelRemoval(name, exceptions, r.matched, r.token),) end function _submodel_removals(::Type{R}, model, prefix) where {R} - prefix = _model_value_varname(model.values, prefix, _model_prefix(model)) + prefix = _model_value_varname(model.values, prefix, getprefix(model)) return mapreduce( r -> _select_removal(r, prefix), (a, b) -> (a..., b...), @@ -4390,7 +4389,7 @@ function _submodel_layer(::Type{R}, model) where {R} prefix = if model.values isa Union{LocalModelValues,UnprefixedArgumentValues} nothing else - _model_prefix(model) + getprefix(model) end return _submodel_values(_binding_layer(R, model.values), prefix) end @@ -4420,7 +4419,7 @@ end # Filter each layer before overlaying: a local fix must not hide an observation # that is recursive in the child namespace. function _submodel_inherited_values(model, prefix) - prefix = _model_value_varname(model.values, prefix, _model_prefix(model)) + prefix = _model_value_varname(model.values, prefix, getprefix(model)) return _inherited_namespace(model.values, prefix) end function _inherited_namespace(values, prefix) diff --git a/src/contexts.jl b/src/contexts.jl index 25e3195e5..d4a6ebc04 100644 --- a/src/contexts.jl +++ b/src/contexts.jl @@ -16,7 +16,7 @@ abstract type AbstractParentContext <: AbstractContext end Return the descendant context of `context`. """ -childcontext +function childcontext end """ setchildcontext(parent::AbstractParentContext, child::AbstractContext) @@ -24,22 +24,8 @@ childcontext Reconstruct `parent` but now using `child` is its [`childcontext`](@ref), effectively updating the child context. -# Examples -```jldoctest; setup=:(using Random) -julia> using DynamicPPL: InitContext, PrefixContext - -julia> ctx = PrefixContext(@varname(a)); - -julia> DynamicPPL.childcontext(ctx) -DefaultContext() - -julia> ctx_prior = DynamicPPL.setchildcontext(ctx, InitContext(MersenneTwister(23), InitFromPrior(), UnlinkAll())); - -julia> DynamicPPL.childcontext(ctx_prior) -InitContext{MersenneTwister, InitFromPrior, UnlinkAll}(MersenneTwister(23), InitFromPrior(), UnlinkAll()) -``` """ -setchildcontext +function setchildcontext end """ leafcontext(context::AbstractContext) diff --git a/src/contexts/prefix.jl b/src/contexts/prefix.jl deleted file mode 100644 index f7376f545..000000000 --- a/src/contexts/prefix.jl +++ /dev/null @@ -1,166 +0,0 @@ -""" - PrefixContext(vn::VarName[, context::AbstractContext]) - PrefixContext(vn::Val{sym}[, context::AbstractContext]) where {sym} - -Create a context that allows you to use the wrapped `context` when running the model and -prefixes all parameters with the VarName `vn`. - -`PrefixContext(Val(:a), context)` is equivalent to `PrefixContext(@varname(a), context)`. -If `context` is not provided, it defaults to `DefaultContext()`. - -This context is useful in nested models to ensure that the names of the parameters are -unique. - -See also: [`to_submodel`](@ref) -""" -struct PrefixContext{Tvn<:VarName,C<:AbstractContext,T} <: AbstractParentContext - vn_prefix::Tvn - context::C - template::T -end -PrefixContext(vn::VarName, context::AbstractContext) = PrefixContext(vn, context, nothing) -PrefixContext(vn::VarName) = PrefixContext(vn, DefaultContext()) -function PrefixContext(::Val{sym}, context::AbstractContext) where {sym} - return PrefixContext(VarName{sym}(), context) -end -PrefixContext(::Val{sym}) where {sym} = PrefixContext(VarName{sym}()) - -childcontext(context::PrefixContext) = context.context -function setchildcontext(ctx::PrefixContext, child::AbstractContext) - return PrefixContext(ctx.vn_prefix, child, ctx.template) -end - -""" - prefix(ctx::AbstractContext, vn::VarName) - -Apply the prefixes in the context `ctx` to the variable name `vn`. -""" -function prefix(ctx::PrefixContext, vn::VarName) - return AbstractPPL.prefix(prefix(childcontext(ctx), vn), ctx.vn_prefix) -end -prefix(::AbstractContext, vn::VarName) = vn -function prefix(ctx::AbstractParentContext, vn::VarName) - return prefix(childcontext(ctx), vn) -end - -""" - optic_skip_length(optic::AbstractPPL.Optic) - -Determine the number of optics that would be added to a VarName by prefixing the VarName -with this optic. - -This is needed when providing templates when setting values in a VarNamedTuple, _within_ a -submodel. That is, inside a submodel, suppose we have `a[1] ~ Normal()`. To call -`templated_setindex!!` for this variable correctly, we need to provide the shape of `a`. Of -course, we can do this, because `a` should be a top-level variable in the model function. -The problem is that `templated_setindex!!` is called from `tilde_assume!!`, which does _not_ -see the variable `a[1]`, but rather `x.a[1]` where `x` is the prefix added by the -`PrefixContext`. - -Thus, to correctly communicate the template, we need to wrap the value for `a` inside N -levels of `SkipTemplate`, which says that "the template is actually for the variable N -levels down". This function computes that N. - -Note that this only counts the prefixes added by `PrefixContext`s; it does not count any -child contexts. That is because each child context will add its own prefixes. -""" -optic_skip_length(::AbstractPPL.Iden) = 0 -optic_skip_length(c::AbstractPPL.Index) = 1 + optic_skip_length(c.child) -optic_skip_length(c::AbstractPPL.Property) = 1 + optic_skip_length(c.child) - -""" - prefix_and_strip_contexts(ctx::PrefixContext, vn::VarName) - -Same as `prefix`, but additionally returns a new context stack that has all the -`PrefixContext`s removed. - -""" -function prefix_and_strip_contexts(ctx::PrefixContext, vn::VarName) - child_context = childcontext(ctx) - # vn_prefixed contains the prefixes from all lower levels - vn_prefixed, child_context_without_prefixes = prefix_and_strip_contexts( - child_context, vn - ) - return AbstractPPL.prefix(vn_prefixed, ctx.vn_prefix), child_context_without_prefixes -end -prefix_and_strip_contexts(ctx::AbstractContext, vn::VarName) = (vn, ctx) -function prefix_and_strip_contexts(ctx::AbstractParentContext, vn::VarName) - vn, new_ctx = prefix_and_strip_contexts(childcontext(ctx), vn) - return vn, setchildcontext(ctx, new_ctx) -end - -""" - DynamicPPL.extract_prefixes(ctx::AbstractContext) - -Recursively process the context stack `ctx` to remove all levels of `PrefixContext`, and -aggregate their prefixes into a single VarName which can be used to apply all the prefixes -at once. - -Returns a tuple of: - -- `ctx_without_prefixes`: the context stack with all `PrefixContext`s removed -- `vn`: the VarName that contains all the aggregated prefixes that were removed, or - `nothing` if there were no `PrefixContext`s in the stack. - -# Example - -```jldoctest -julia> using DynamicPPL: extract_prefixes, PrefixContext, AbstractContext, @varname - -julia> struct MyContext <: AbstractContext end - -julia> ctx = PrefixContext(@varname(b), PrefixContext(@varname(a), MyContext())); - -julia> extract_prefixes(ctx) -(MyContext(), b.a) -``` - -""" -function extract_prefixes(ctx::AbstractContext) - new_ctx, new_optic = _extract_prefixes(ctx, AbstractPPL.Iden()) - return if new_optic === AbstractPPL.Iden() - new_ctx, nothing - else - new_ctx, AbstractPPL.optic_to_varname(new_optic) - end -end -function _extract_prefixes(ctx::PrefixContext, optic::AbstractPPL.AbstractOptic) - new_child_context, child_optic = _extract_prefixes(childcontext(ctx), optic) - this_optic = AbstractPPL.varname_to_optic(ctx.vn_prefix) - return (new_child_context, Base.cat(this_optic, child_optic)) -end -function _extract_prefixes(ctx::AbstractParentContext, optic::AbstractPPL.AbstractOptic) - # ctx is not a PrefixContext, so just need to recurse on the child context - new_child_context, optic = _extract_prefixes(childcontext(ctx), optic) - return setchildcontext(ctx, new_child_context), optic -end -function _extract_prefixes(ctx::AbstractContext, optic::AbstractPPL.AbstractOptic) - return ctx, optic -end - -function tilde_assume!!( - context::PrefixContext, - right::Distribution, - vn::VarName, - template::Any, - vi::AbstractVarInfo, -) - # Note that we can't use something like this here: - # new_vn = prefix(context, vn) - # return tilde_assume!!(childcontext(context), right, new_vn, vi) - # This is because `prefix` applies _all_ prefixes in a given context to a - # variable name. Thus, if we had two levels of nested prefixes e.g. - # `PrefixContext{:a}(PrefixContext{:b}(DefaultContext()))`, then the - # first call would apply the prefix `a.b._`, and the recursive call - # would apply the prefix `b._`, resulting in `b.a.b._`. - # This is why we need a special function, `extract_prefixes`. - new_context, aggregated_prefix_vn = extract_prefixes(context) - new_vn = AbstractPPL.prefix(vn, aggregated_prefix_vn) - # Figure out how many levels of templating to skip for this variable (and add 1 for the - # top-level symbol in the VarName). - # NOTE(penelopeysm): I tried to move this computation into an inner constructor of - # PrefixContext. But that caused nonzero allocations on the LogDensityFunction submodel - # test, for reasons that are rather unclear! Be careful if you think of doing that. - n = optic_skip_length(AbstractPPL.getoptic(aggregated_prefix_vn)) + 1 - return tilde_assume!!(new_context, right, new_vn, SkipTemplate{n}(template), vi) -end diff --git a/src/debug_utils.jl b/src/debug_utils.jl index 326608ec1..1ee8dae60 100644 --- a/src/debug_utils.jl +++ b/src/debug_utils.jl @@ -58,19 +58,19 @@ function DynamicPPL.tilde_assume!!( _register_removals!(ctx, submodel.model) local_prefix = if AutoPrefix namespace = DynamicPPL._concretize_prefix(vn, template; prefix=Val(true)) - DynamicPPL.maybe_prefix(DynamicPPL._model_prefix(submodel.model), namespace) + DynamicPPL.maybe_prefix(DynamicPPL.getprefix(submodel.model), namespace) else - DynamicPPL._model_prefix(submodel.model) + DynamicPPL.getprefix(submodel.model) end - full_prefix = DynamicPPL.maybe_prefix(local_prefix, DynamicPPL._model_prefix(parent)) + full_prefix = DynamicPPL.maybe_prefix(local_prefix, DynamicPPL.getprefix(parent)) lock(ctx.lock) do push!(ctx.template_models, (submodel.model, full_prefix)) end - prefixed = AutoPrefix || DynamicPPL._model_prefix(submodel.model) !== nothing + prefixed = AutoPrefix || DynamicPPL.getprefix(submodel.model) !== nothing if ctx.track_names lock(ctx.lock) do if prefixed - namespace = AutoPrefix ? vn : DynamicPPL._model_prefix(submodel.model) + namespace = AutoPrefix ? vn : DynamicPPL.getprefix(submodel.model) push!(ctx.namespaces, DynamicPPL.AbstractPPL.getsym(namespace)) else push!(ctx.models, submodel.model) @@ -333,7 +333,7 @@ function check_model( model.context, Model[model], Set{Symbol}(), ReentrantLock() ) _register_removals!(binding_context, model) - push!(binding_context.template_models, (model, DynamicPPL._model_prefix(model))) + push!(binding_context.template_models, (model, DynamicPPL.getprefix(model))) checked_model = DynamicPPL.contextualize(model, binding_context) _, vi = DynamicPPL.init!!(rng, checked_model, vi, init_strategy, UnlinkAll()) _warn_unused_binding_names(model, binding_context) diff --git a/src/model.jl b/src/model.jl index 3078ae88f..1a5f82939 100644 --- a/src/model.jl +++ b/src/model.jl @@ -117,6 +117,9 @@ function _apply_prefix_template(prefix::PrefixTemplate, template) ) end _compose_prefix_templates(::Nothing, inner) = inner +function _compose_prefix_templates(prefix::VarName, inner::Union{Nothing,VarName}) + return maybe_prefix(inner, prefix) +end function _compose_prefix_templates(prefix::VarName, inner) return PrefixTemplate(prefix, NoTemplate(), inner) end @@ -126,6 +129,9 @@ function _compose_prefix_templates(prefix::PrefixTemplate, inner) ) end +_getprefix(prefix::Union{Nothing,VarName}) = prefix +_getprefix(prefix::PrefixTemplate) = maybe_prefix(_getprefix(prefix.inner), prefix.prefix) + is_splat_symbol(s::Symbol) = startswith(string(s), "#splat#") function unsplat_symbol(s::Symbol) return is_splat_symbol(s) ? Symbol(chopprefix(string(s), "#splat#")) : s @@ -219,10 +225,11 @@ struct Model{ defaultnames, Targs, Tdefaults, - C<:AbstractContext, + Prefix<:Union{VarName,Nothing,PrefixTemplate}, Values<:Union{ VarNamedTuple,ModelBindingLayers,LocalModelValues,UnprefixedArgumentValues }, + C<:AbstractContext, Threaded, ArgsOnLHS, } <: AbstractProbabilisticProgram @@ -230,15 +237,17 @@ struct Model{ args::NamedTuple{argnames,Targs} defaults::NamedTuple{defaultnames,Tdefaults} context::C + prefix::Prefix values::Values function Model{Threaded}( f::F, args::NamedTuple{A,Ta}, defaults::NamedTuple{D,Td}, - context::C, - values::V; + prefix::P, + values::V, + context::C; args_on_lhs::Union{Tuple{Vararg{Symbol}},Vector{Symbol},ModelBindingMetadata}=(), - ) where {F,A,Ta,D,Td,C,V,Threaded} + ) where {F,A,Ta,D,Td,P,C,V,Threaded} mapreduce( pair -> pair.second isa ModelValue, &, _model_values(values); init=true ) || throw(ArgumentError("Model values must carry a condition or fix role")) @@ -251,21 +260,31 @@ struct Model{ ), ) end - return new{F,A,D,Ta,Td,C,V,Threaded,metadata}(f, args, defaults, context, values) + return new{F,A,D,Ta,Td,P,V,C,Threaded,metadata}( + f, args, defaults, context, prefix, values + ) end # Internal reconstruction reuses already-validated bindings. function DynamicPPL._reconstruct_model( - model::Model{F,A,D,Ta,Td}, context::C, values::V, ::Val{Threaded} - ) where {F,A,D,Ta,Td,C,V,Threaded} - return new{F,A,D,Ta,Td,C,V,Threaded,_binding_metadata(model)}( - model.f, model.args, model.defaults, context, values + model::Model{F,A,D,Ta,Td}, prefix::P, context::C, values::V, ::Val{Threaded} + ) where {F,A,D,Ta,Td,P,C,V,Threaded} + return new{F,A,D,Ta,Td,P,V,C,Threaded,_binding_metadata(model)}( + model.f, model.args, model.defaults, context, prefix, values ) end end +""" + getprefix(model::Model) + +Return the combined prefix of the model's LHS variables, or `nothing` when absent. +Storage templates for nested submodel namespaces remain internal to the model. +""" +getprefix(model::Model) = _getprefix(model.prefix) + function _binding_metadata( - ::Model{F,A,D,Ta,Td,C,V,Threaded,ArgsOnLHS} -) where {F,A,D,Ta,Td,C,V,Threaded,ArgsOnLHS} + ::Model{F,A,D,Ta,Td,P,V,C,Threaded,ArgsOnLHS} +) where {F,A,D,Ta,Td,P,V,C,Threaded,ArgsOnLHS} return ArgsOnLHS end @@ -279,7 +298,7 @@ Base.@constprop :aggressive function Model{Threaded}( args_on_lhs::Union{Tuple{Vararg{Symbol}},Vector{Symbol},ModelBindingMetadata}=(), ) where {Threaded} values = _argument_defaults(merge(args, defaults), Val(_args_on_lhs(args_on_lhs))) - model = Model{Threaded}(f, args, defaults, context, values; args_on_lhs) + model = Model{Threaded}(f, args, defaults, nothing, values, context; args_on_lhs) found = _placeholder_names(merge(args, defaults), Val(_args_on_lhs(args_on_lhs))) return all(isempty, found) ? model : _remove_placeholders(model, found) end @@ -347,10 +366,15 @@ end Return whether `model` has been marked as needing threadsafe evaluation (using `setthreadsafe`). """ -requires_threadsafe(::Model{F,A,D,Ta,Td,C,V,Threaded}) where {F,A,D,Ta,Td,C,V,Threaded} = - Threaded -function _reconstruct_model(model::Model; context=model.context, values=model.values) - return _reconstruct_model(model, context, values, Val(requires_threadsafe(model))) +requires_threadsafe( + ::Model{F,A,D,Ta,Td,P,V,C,Threaded} +) where {F,A,D,Ta,Td,P,V,C,Threaded} = Threaded +function _reconstruct_model( + model::Model; prefix=model.prefix, context=model.context, values=model.values +) + return _reconstruct_model( + model, prefix, context, values, Val(requires_threadsafe(model)) + ) end """ @@ -359,23 +383,12 @@ end Return a model with its context replaced by `context`. """ function contextualize(model::Model, context::AbstractContext) - if model.values isa UnprefixedArgumentValues && ( - last(extract_prefixes(context)) != _model_prefix(model) || - _prefix_template(context) !== _model_prefix_template(model) - ) - model = _materialize_argument_values(model) - end return _reconstruct_model(model; context) end """Return a model with its leaf context replaced by `context`.""" function setleafcontext(model::Model, context::AbstractContext) return contextualize(model, setleafcontext(model.context, context)) end -_model_prefix(model::Model) = last(extract_prefixes(model.context)) -_prefix_template(::AbstractContext) = nothing -_prefix_template(context::AbstractParentContext) = _prefix_template(childcontext(context)) -_prefix_template(context::PrefixContext) = context.template -_model_prefix_template(model::Model) = _prefix_template(model.context) """ setthreadsafe(model::Model, threadsafe::Bool) @@ -400,7 +413,9 @@ function setthreadsafe(model::Model, threadsafe::Bool) return if requires_threadsafe(model) == threadsafe model else - _reconstruct_model(model, model.context, model.values, Val(threadsafe)) + _reconstruct_model( + model, model.prefix, model.context, model.values, Val(threadsafe) + ) end end @@ -539,7 +554,7 @@ function prefix(model::Model, x::VarName; template=NoTemplate()) model = _materialize_argument_values(model) values = if model.values isa VarNamedTuple && - _model_prefix(model) === nothing && + getprefix(model) === nothing && !isempty(model.values) && mapreduce( pair -> pair.second isa ModelValue{ArgumentCondition}, @@ -554,22 +569,12 @@ function prefix(model::Model, x::VarName; template=NoTemplate()) return _prefix_model(model, x, template, values) end function _prefix_model(model::Model, x::VarName, template, values) - model_prefix = maybe_prefix(_model_prefix(model), x) - prefix_template = - if template isa NoTemplate && _model_prefix_template(model) === nothing - nothing - else - inner = if _model_prefix_template(model) === nothing - _model_prefix(model) - else - _model_prefix_template(model) - end - PrefixTemplate(x, template, inner) - end - context = PrefixContext( - model_prefix, first(extract_prefixes(model.context)), prefix_template - ) - return _reconstruct_model(model; context, values) + prefix = if template isa NoTemplate + _compose_prefix_templates(x, model.prefix) + else + PrefixTemplate(x, template, model.prefix) + end + return _reconstruct_model(model; prefix, values) end function prefix(model::Model, ::Val{sym}) where {sym} return prefix(model, VarName{sym}()) @@ -579,14 +584,13 @@ function prefix(model::Model, x) end function _prefix_varname_and_template(vn::VarName, template::Any, model::Model) - return _prefix_varname_and_template( - vn, template, _model_prefix(model), _model_prefix_template(model) - ) + return _prefix_varname_and_template(vn, template, model.prefix) end -function _prefix_varname_and_template(vn::VarName, template, prefix, prefix_template) +function _prefix_varname_and_template(vn::VarName, template, prefix) prefix === nothing && return vn, template - pt = prefix_template === nothing ? prefix : prefix_template - return AbstractPPL.prefix(vn, prefix), _apply_prefix_template(pt, template) + return ( + AbstractPPL.prefix(vn, _getprefix(prefix)), _apply_prefix_template(prefix, template) + ) end function tilde_assume!!( @@ -623,7 +627,7 @@ function _check_tilde_value(value, vn, ::Union{Condition,Fix}) end """ - tilde_observe!!(prefix, prefix_template, right::Distribution, left, vn, template, vi) + tilde_observe!!(prefix, right::Distribution, left, vn, template, vi) Accumulate an observation and return `(left, vi)` with the updated varinfo. @@ -631,18 +635,17 @@ Accumulate an observation and return `(left, vi)` with the updated varinfo. the variable name before prefixing, or `nothing` for a literal. `template` describes the top-level variable's storage; literals use `NoTemplate()`. -Apply `prefix` (a `VarName` or `nothing`) and its storage template `prefix_template`, -then delegate to [`accumulate_observe!!`](@ref). The compiler passes this metadata directly -so observations do not box the model. Every observation calls this function, independently +Apply `prefix`, the model's single prefix value (`nothing`, a `VarName`, or a +`PrefixTemplate` containing storage metadata), then delegate to [`accumulate_observe!!`](@ref). +The compiler passes this metadata directly so observations do not box the model. +Every observation calls this function, independently of the evaluation context. Fixed LHS variables bypass it and do not contribute to the log probability. """ -function tilde_observe!!( - prefix, prefix_template, right::Distribution, left, vn, template, vi -) +function tilde_observe!!(prefix, right::Distribution, left, vn, template, vi) vn, template = if vn === nothing vn, NoTemplate() else - _prefix_varname_and_template(vn, template, prefix, prefix_template) + _prefix_varname_and_template(vn, template, prefix) end left = _check_tilde_value(left, vn, Condition()) vi = accumulate_observe!!(vi, right, _nonmissing_eltype(left), vn, template) @@ -1104,3 +1107,7 @@ end function AbstractPPL.evaluate!!(model::Model, context::AbstractContext, vi::AbstractVarInfo) return evaluate_nowarn!!(setleafcontext(model, context), vi) end + +optic_skip_length(::AbstractPPL.Iden) = 0 +optic_skip_length(optic::AbstractPPL.Index) = 1 + optic_skip_length(optic.child) +optic_skip_length(optic::AbstractPPL.Property) = 1 + optic_skip_length(optic.child) diff --git a/src/submodel.jl b/src/submodel.jl index 95d14492e..0cad1a5d4 100644 --- a/src/submodel.jl +++ b/src/submodel.jl @@ -183,7 +183,7 @@ function _submodel_namespace(::Union{ModelValue,ModelValueTree}) end function _submodel_values(model::Model, prefix) - prefix = _model_value_varname(model.values, prefix, _model_prefix(model)) + prefix = _model_value_varname(model.values, prefix, getprefix(model)) return _submodel_values(_model_values(model.values), prefix) end _submodel_values(values::VarNamedTuple, ::Nothing) = values @@ -195,7 +195,7 @@ end # Shape ownership survives selecting a child's namespace, including whole namespace bindings. function _submodel_fixed_owners(model::Model, prefix) - prefix = _model_value_varname(model.values, prefix, _model_prefix(model)) + prefix = _model_value_varname(model.values, prefix, getprefix(model)) owners = _fixed_owners(model.values) prefix === nothing && return owners return mapreduce((a, b) -> (a..., b...), owners; init=()) do owner @@ -241,13 +241,13 @@ Evaluate `submodel` under `parent_model`. ) end !AutoPrefix && - _model_prefix(submodel.model) === nothing && + getprefix(submodel.model) === nothing && _check_shared_removals(parent_model, submodel.model) left_vn = AutoPrefix ? _concretize_prefix(left_vn, template; prefix=Val(true)) : left_vn local_prefix = if AutoPrefix - maybe_prefix(_model_prefix(submodel.model), left_vn) + maybe_prefix(getprefix(submodel.model), left_vn) else - _model_prefix(submodel.model) + getprefix(submodel.model) end inherited_templates = _submodel_binding_templates(parent_model, local_prefix) observation_removals = _submodel_removals(Condition, parent_model, local_prefix) @@ -275,10 +275,10 @@ Evaluate `submodel` under `parent_model`. _submodel_inherited_values(parent_model, local_prefix), inherited_templates, ), - maybe_prefix(local_prefix, _model_prefix(parent_model)), + maybe_prefix(local_prefix, getprefix(parent_model)), ) # Shared unprefixed names may belong to the parent or another child. - if AutoPrefix || _model_prefix(submodel.model) !== nothing + if AutoPrefix || getprefix(submodel.model) !== nothing _check_binding_addresses(child_model, parent_values) end owners = (child_owners..., _submodel_fixed_owners(parent_model, local_prefix)...) @@ -287,7 +287,7 @@ Evaluate `submodel` under `parent_model`. _merge_model_values( child_values, parent_values, - maybe_prefix(local_prefix, _model_prefix(parent_model)), + maybe_prefix(local_prefix, getprefix(parent_model)), ), ( _submodel_removals(Condition, submodel.model, nothing)..., @@ -319,7 +319,7 @@ Evaluate `submodel` under `parent_model`. isempty(observation_removals) && isempty(fixed_removals) || _check_shared_latent_storage( _reconstruct_model(submodel.model; values), - maybe_prefix(local_prefix, _model_prefix(parent_model)), + maybe_prefix(local_prefix, getprefix(parent_model)), ) return _evaluate_submodel!!( parent_model, context, submodel, left_vn, template, vi, values @@ -336,37 +336,12 @@ end vi::AbstractVarInfo, values::LocalModelValues, ) where {M,AutoPrefix} - parent_prefix = _model_prefix(parent_model) model = if AutoPrefix vn, template = _prefix_varname_and_template(left_vn, template, parent_model) _prefix_model(submodel.model, vn, template, values) - elseif parent_prefix === nothing - _reconstruct_model(submodel.model; values) else - model = _prefix_model( - submodel.model, - parent_prefix, - _apply_prefix_template(_model_prefix_template(parent_model), NoTemplate()), - values, - ) - if _model_prefix_template(parent_model) === nothing - model - else - inner = if _model_prefix_template(submodel.model) === nothing - _model_prefix(submodel.model) - else - _model_prefix_template(submodel.model) - end - prefix_template = _compose_prefix_templates( - _model_prefix_template(parent_model), inner - ) - prefix_context = PrefixContext( - _model_prefix(model), - first(extract_prefixes(model.context)), - prefix_template, - ) - _reconstruct_model(model; context=prefix_context) - end + prefix = _compose_prefix_templates(parent_model.prefix, submodel.model.prefix) + _reconstruct_model(submodel.model; prefix, values) end # Calling model.f directly avoids the inference recursion limit as nested prefixes # change the Model type; routing through _evaluate!! widens it to Any (Turing.jl#2844). @@ -375,8 +350,6 @@ end return model.f(args...; kwargs...) end -function tilde_observe!!( - prefix, prefix_template, ::DynamicPPL.Submodel, left, ::Nothing, template, vi -) +function tilde_observe!!(prefix, ::DynamicPPL.Submodel, left, ::Nothing, template, vi) throw(ArgumentError("`x ~ to_submodel(...)` is not supported when `x` is a literal")) end diff --git a/test/compiler.jl b/test/compiler.jl index fb4d40907..45515bc30 100644 --- a/test/compiler.jl +++ b/test/compiler.jl @@ -1020,8 +1020,12 @@ end @testset "public observation call target" begin count_observe(::Any) = 0 function count_observe(ex::Expr) - return (Meta.isexpr(ex, :call) && ex.args[1] === tilde_observe!!) + - sum(count_observe, ex.args; init=0) + is_observe = Meta.isexpr(ex, :call) && ex.args[1] === tilde_observe!! + if is_observe + @test length(ex.args) == 7 # Function and six arguments. + @test ex.args[2] == :(__model__.prefix) + end + return is_observe + sum(count_observe, ex.args; init=0) end expr = @macroexpand @model function observations(x) x ~ Normal() diff --git a/test/context_implementations.jl b/test/context_implementations.jl index ba4dbe04e..e7b793d5a 100644 --- a/test/context_implementations.jl +++ b/test/context_implementations.jl @@ -15,13 +15,15 @@ struct UnimplementedStrategy <: AbstractInitStrategy end struct ObserveHookDistribution <: ContinuousUnivariateDistribution seen::Vector{Union{VarName,Nothing}} + prefixes::Vector{Any} end Distributions.logpdf(::ObserveHookDistribution, x::Real) = logpdf(Normal(), x) function DynamicPPL.tilde_observe!!( - prefix, prefix_template, dist::ObserveHookDistribution, left, vn, template, vi + prefix, dist::ObserveHookDistribution, left, vn, template, vi ) push!(dist.seen, vn) - return tilde_observe!!(prefix, prefix_template, Normal(), left, vn, template, vi) + push!(dist.prefixes, prefix) + return tilde_observe!!(prefix, Normal(), left, vn, template, vi) end struct RecordingStrategy{S<:AbstractInitStrategy} <: AbstractInitStrategy @@ -57,15 +59,47 @@ end end seen = Union{VarName,Nothing}[] model = condition( - observations(1.0, [2.0, 3.0], ObserveHookDistribution(seen)); z=4.0 + observations(1.0, [2.0, 3.0], ObserveHookDistribution(seen, Any[])); z=4.0 ) - _, vi = DynamicPPL.evaluate_nowarn!!(model, VarInfo()) - @test seen == [@varname(x), @varname(z), nothing, @varname(ys[1]), @varname(ys[2])] - @test getloglikelihood(vi) ≈ sum(logpdf.(Normal(), 0.0:4.0)) - empty!(seen) - _, vi = DynamicPPL.evaluate_nowarn!!(outer(model), VarInfo()) - @test length(seen) == 5 - @test getloglikelihood(vi) ≈ sum(logpdf.(Normal(), 0.0:4.0)) + templated = prefix(model, @varname(a[2, 3]); template=zeros(2, 3)) + nested = outer(templated) + for m in (model, prefix(model, @varname(p)), templated, nested) + empty!(seen) + prefixes = model.args.dist.prefixes + empty!(prefixes) + acc = DynamicPPL.VNTAccumulator{DynamicPPL.POINTWISE_ACCNAME}( + DynamicPPL.PointwiseLogProb{false,true}() + ) + _, vi = DynamicPPL.evaluate_nowarn!!( + m, VarInfo((DynamicPPL.default_accumulators()..., acc)) + ) + @test seen == + [@varname(x), @varname(z), nothing, @varname(ys[1]), @varname(ys[2])] + @test length(prefixes) == 5 + if m === nested + @test all(p -> p isa DynamicPPL.PrefixTemplate, prefixes) + @test all(p -> DynamicPPL._getprefix(p) == @varname(b.a[2, 3]), prefixes) + else + @test all(p -> p === m.prefix, prefixes) + end + @test getloglikelihood(vi) ≈ sum(logpdf.(Normal(), 0.0:4.0)) + pointwise = DynamicPPL.get_pointwise_logprobs(vi) + observation_prefix = DynamicPPL._getprefix(first(prefixes)) + for (vn, value) in ( + (@varname(x), 1.0), + (@varname(z), 4.0), + (@varname(ys[1]), 2.0), + (@varname(ys[2]), 3.0), + ) + name = DynamicPPL.maybe_prefix(vn, observation_prefix) + @test pointwise[name] ≈ logpdf(Normal(), value) + end + if m === templated + @test size(pointwise.data.a.data) == (2, 3) + elseif m === nested + @test size(pointwise.data.b.data.a.data) == (2, 3) + end + end end @testset "prefixed evaluation: $(model.f)" for model in DynamicPPL.TestUtils.ALL_MODELS @@ -105,7 +139,7 @@ end for value in (one(T), T[1, 2]) for vn in (@varname(x), nothing) result, vi = @inferred tilde_observe!!( - nothing, NoTemplate(), dist, value, vn, NoTemplate(), VarInfo() + nothing, dist, value, vn, NoTemplate(), VarInfo() ) @test result === value @test getloglikelihood(vi) ≈ loglikelihood(dist, value) diff --git a/test/debug_utils.jl b/test/debug_utils.jl index 9afdee1b8..d95598aa4 100644 --- a/test/debug_utils.jl +++ b/test/debug_utils.jl @@ -180,8 +180,9 @@ end __model__.f, (; y, checking=true), __model__.defaults, - __model__.context, - __model__.values; + __model__.prefix, + __model__.values, + __model__.context; args_on_lhs=DynamicPPL._args_on_lhs(__model__), ) @test check_model(child) diff --git a/test/model.jl b/test/model.jl index a9cbb76fc..dd17e0382 100644 --- a/test/model.jl +++ b/test/model.jl @@ -129,7 +129,9 @@ DynamicPPL.convert_model_argument(T, ::Type{CustomModelArgument}) = (:converted_ for m in ( DynamicPPL.Model{false}(f, (; x=missing, y=1.0), (;)), DynamicPPL.Model{false}(f, (; x=missing); y=1.0), - DynamicPPL.Model{false}(f, (; x=missing, y=1.0), (;), DefaultContext()), + DynamicPPL.Model{false}( + f, (; x=missing, y=1.0), (;), nothing, VarNamedTuple(), DefaultContext() + ), ) @test isempty(conditioned(m)) @test isempty(DynamicPPL._args_on_lhs(m)) @@ -157,6 +159,31 @@ DynamicPPL.convert_model_argument(T, ::Type{CustomModelArgument}) = (:converted_ end end + @testset "direct construction preserves merged prefix metadata" begin + @model prefixed_lhs(y) = y ~ Normal() + original = prefixed_lhs(1.0) + for model in ( + original, + prefix(original, @varname(p)), + prefix(original, @varname(p[2]); template=zeros(2)), + ) + direct = DynamicPPL.Model{false}( + model.f, + model.args, + model.defaults, + model.prefix, + model.values, + model.context; + args_on_lhs=DynamicPPL._args_on_lhs(model), + ) + @test direct.prefix === model.prefix + @test (@inferred DynamicPPL.getprefix(direct)) == DynamicPPL.getprefix(model) + @test conditioned(direct) == conditioned(model) + @test keys(VarInfo(Xoshiro(1), decondition(direct))) == + keys(VarInfo(Xoshiro(1), decondition(model))) + end + end + @testset "direct construction records argument-supplied observations" begin @model positional_lhs(y, mu) = y ~ Normal(mu) @model keyword_lhs(; y, mu) = y ~ Normal(mu) diff --git a/test/prefix.jl b/test/prefix.jl new file mode 100644 index 000000000..22d15c69c --- /dev/null +++ b/test/prefix.jl @@ -0,0 +1,124 @@ +module DynamicPPLPrefixTests + +using AbstractPPL: AbstractPPL +using Dates: now +using Distributions +using DynamicPPL +using Test +using Random: Xoshiro + +@info "Testing $(@__FILE__)..." +__now__ = now() + +@testset "model prefixes" begin + @model function demo_prefix() + x ~ Normal() + y = Vector{Float64}(undef, 2) + y[1] ~ Normal() + y[2] ~ Normal() + return (; x, y) + end + + model = demo_prefix() + prefixed = @inferred DynamicPPL.prefix(model, @varname(a)) + twice_prefixed = @inferred DynamicPPL.prefix(prefixed, @varname(b)) + + @test Set(keys(rand(prefixed))) == + Set([@varname(a.x), @varname(a.y[1]), @varname(a.y[2])]) + @test Set(keys(rand(twice_prefixed))) == + Set([@varname(b.a.x), @varname(b.a.y[1]), @varname(b.a.y[2])]) + + @testset "compound prefixes" begin + for prefix_vn in (@varname(a.b), @varname(a[1])) + transformed = DynamicPPL.prefix(model, prefix_vn) + @test Set(keys(rand(transformed))) == Set([ + AbstractPPL.prefix(@varname(x), prefix_vn), + AbstractPPL.prefix(@varname(y[1]), prefix_vn), + AbstractPPL.prefix(@varname(y[2]), prefix_vn), + ]) + end + end + + @testset "stored values follow operation order" begin + conditioned_first = DynamicPPL.prefix(condition(model; x=1.0), @varname(a)) + @test conditioned(conditioned_first)[@varname(a.x)] == 1.0 + @test conditioned_first().x == 1.0 + + prefixed_first = condition( + DynamicPPL.prefix(model, @varname(a)), @varname(a.x) => 2.0 + ) + @test conditioned(prefixed_first)[@varname(a.x)] == 2.0 + @test prefixed_first().x == 2.0 + + fixed_first = DynamicPPL.prefix(fix(model; x=3.0), @varname(a)) + @test fixed(fixed_first)[@varname(a.x)] == 3.0 + @test fixed_first().x == 3.0 + + prefixed_then_fixed = fix( + DynamicPPL.prefix(model, @varname(a)), @varname(a.x) => 4.0 + ) + @test fixed(prefixed_then_fixed)[@varname(a.x)] == 4.0 + @test prefixed_then_fixed().x == 4.0 + + @test isempty(keys(conditioned(decondition(conditioned_first, @varname(a))))) + @test isempty(keys(fixed(unfix(fixed_first, @varname(a))))) + end + + @testset "explicit bindings stay under the model prefix" begin + @model argument_lhs(y) = y ~ Normal() + for op in (condition, fix) + prefixed_argument = DynamicPPL.prefix(argument_lhs(1.0), @varname(p)) + message = "Cannot bind `y`: it is outside this model's prefix `p`." + for bindings in ((; y=0.0), @varname(y) => 0.0, VarNamedTuple(; y=0.0)) + @test_throws message op(prefixed_argument, bindings) + end + @test op(prefixed_argument, @varname(p.y) => 2.0)(Xoshiro(1)) == 2.0 + @test op(prefixed_argument; p=(y=3.0,))(Xoshiro(1)) == 3.0 + @test_throws "not an LHS top symbol" op( + prefixed_argument, @varname(p.unused) => 0.0 + ) + + nested = DynamicPPL.prefix(prefixed_argument, @varname(q)) + @test_throws ArgumentError op(nested, @varname(p.y) => 0.0) + @test_throws ArgumentError op(nested, @varname(q.other.y) => 0.0) + @test op(nested, @varname(q.p.y) => 4.0)(Xoshiro(1)) == 4.0 + + indexed = DynamicPPL.prefix(argument_lhs(1.0), @varname(p[1])) + @test_throws ArgumentError op(indexed, @varname(p[2].y) => 0.0) + @test op(indexed, @varname(p[1].y) => 5.0)(Xoshiro(1)) == 5.0 + + @model function parent_binding(child) + result ~ to_submodel(child) + return result + end + @model function parent_binding_unprefixed(child) + result ~ to_submodel(child, false) + return result + end + for auto_prefix in (false, true) + child = op(nested, @varname(q.p.y) => 6.0) + parent = + auto_prefix ? parent_binding(child) : parent_binding_unprefixed(child) + @test parent(Xoshiro(1)) == 6.0 + address = auto_prefix ? @varname(result.q.p.y) : @varname(q.p.y) + @test op(parent, address => 7.0)(Xoshiro(1)) == 7.0 + end + end + end + + @testset "dynamic prefixes use the supplied template" begin + for op in (condition, fix) + prefixed = @inferred DynamicPPL.prefix( + op(model; x=2.0), @varname(a[end, end]); template=zeros(2, 3) + ) + @test prefixed().x == 2.0 + values = op === condition ? conditioned(prefixed) : fixed(prefixed) + @test size(values.data.a.data) == (2, 3) + @test values[@varname(a[2, 3].x)] == 2.0 + end + end +end + +@info "Completed $(@__FILE__) in $(now() - __now__)." + +end diff --git a/test/runtests.jl b/test/runtests.jl index 29c9ca8e5..74947f594 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -38,6 +38,7 @@ Random.seed!(100) include("lkj.jl") include("contexts/init.jl") include("conditionfix.jl") + include("prefix.jl") include("context_implementations.jl") include("threadsafe.jl") include("debug_utils.jl") diff --git a/test/submodels.jl b/test/submodels.jl index 3275b603e..10bfbba86 100644 --- a/test/submodels.jl +++ b/test/submodels.jl @@ -1344,7 +1344,7 @@ end model = prefix( prefix_leaf(), @varname(p[2].q[end]); template=[(q=zeros(2),), (q=zeros(2),)] ) - @test DynamicPPL._model_prefix(model) == @varname(p[2].q[2]) + @test DynamicPPL.getprefix(model) == @varname(p[2].q[2]) @test model(Xoshiro(1)) == 3.0 end for (name, canonical) in ( @@ -1356,7 +1356,7 @@ end ) for leaf in (prefix_leaf(), condition(prefix_leaf(); x=3.0)) model = prefix(leaf, name; template=zeros(2, 2)) - @test DynamicPPL._model_prefix(model) == canonical + @test DynamicPPL.getprefix(model) == canonical @test model(Xoshiro(1)) == 3.0 @test prefix_nested(model)(Xoshiro(1)) == 3.0 end