Migrating old VarInfo code

OnlyAccsVarInfo is removed. Use VarInfo, which has the same accumulator constructor forms: VarInfo(accs...), VarInfo(accs::Tuple), and VarInfo(accs::AccumulatorTuple). The old VarInfo{Tfm,T,Accs} is now VarInfo{Accs}; update dispatch that uses the old type parameters. Transform strategies are evaluation inputs, not part of the output type. VarInfo() records log densities, not parameter values. Add a RawValueAccumulator 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.

To reuse previous values, extract them explicitly before evaluating:

context = InitContext(rng, InitFromParams(get_vector_values(previous), nothing), LinkAll())
retval, outputs = evaluate!!(model, context, VarInfo())

Use get_raw_values(previous) for a raw-value accumulator. Choose UnlinkAll(), LinkAll(), or a partial/fixed transform strategy for the desired output representation. init!! previously inferred its default transform strategy from the VarInfo. It now defaults to UnlinkAll() regardless of the recorded values. Replace init!!(rng, model, vi, init) with init!!(rng, model, vi, init, strategy) when the outputs should use a different representation. For custom strategies, implement get_param_eltype(strategy) when evaluation needs to promote argument buffers or thread-local accumulators for AD.

All output accumulators reset before evaluation. Previously recorded sites that are not executed again are absent from the new outputs. Replace vi.values with get_vector_values(vi); there is no separate parameter store.

Please get in touch if you have some old code you're unsure how to migrate, and we will be happy to add it to this list.

using DynamicPPL, Distributions, Random
using BangBang: BangBang

@model function f()
    x ~ Normal()
    y ~ LogNormal()
    return 1.0 ~ Normal(x + y)
end

model = f()

Model{typeof(Main.f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false}(Main.f, NamedTuple(), NamedTuple(), DefaultContext())

Sampling from the prior

Old:

vi = VarInfo(Xoshiro(468), model)

New:

accs = VarInfo()
_, vi = init!!(Xoshiro(468), model, accs, InitFromPrior(), UnlinkAll())
vi

VarInfo
 └─ AccumulatorTuple with 3 accumulators
    ├─ LogPrior => LogPriorAccumulator(-1.769167187291674)
    ├─ LogJacobian => LogJacobianAccumulator(0.0)
    └─ LogLikelihood => LogLikelihoodAccumulator(-0.918938737957795)

Getting parameter values

Old:

vi = VarInfo(Xoshiro(468), model)
vi.values[@varname(x)]

New:

# Set to true if you want to include results of `:=` statements.
accs = VarInfo(RawValueAccumulator(false))
_, vi = init!!(Xoshiro(468), model, accs, InitFromPrior(), UnlinkAll())
get_raw_values(vi)

VarNamedTuple
├─ x => 0.07200886749732076
└─ y => 0.9286310592520649

Generating vectorised parameters from linked VarInfo

Old:

vi = VarInfo(Xoshiro(468), model)
vi = DynamicPPL.link!!(vi, model)
vi[:]

The new pattern recognises that in practice you are likely using vi[:] in conjunction with a LogDensityFunction. So we make one first:

ldf = LogDensityFunction(model, getlogjoint_internal, LinkAll())

Then you can do:

rand(Xoshiro(468), ldf)

2-element Vector{Float64}:
  0.07200886749732076
 -0.07404375655951738

This gives you a set of parameters, but if you want to also obtain the log-density at the new parameters, you can do this in a single call to init!!; please see the documentation on LogDensityFunction for more details on how to do this.

Re-evaluating log density at new parameters

Old:

vi = VarInfo(Xoshiro(468), model)

vals = [1.0, 1.0]
vi = DynamicPPL.unflatten!!(vi, vals)
_, vi = DynamicPPL.evaluate!!(model, vi)
vi

The new path also assumes that you are using a LogDensityFunction:

# Note that we use `UnlinkAll()` here to match the VarInfo above.
# If your VarInfo was linked, you should use `LinkAll()` instead.

ldf = LogDensityFunction(model, getlogjoint_internal, UnlinkAll())

LogDensityFunction{Model{typeof(Main.f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false}, Nothing, UnlinkAll, typeof(getlogjoint_internal), VarNamedTuple{(:x, :y), Tuple{RangeAndTransform{Unlink}, RangeAndTransform{Unlink}}}, Nothing, Vector{Float64}, DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}, false}(Model{typeof(Main.f), (), (), (), Tuple{}, Tuple{}, DefaultContext, false}(Main.f, NamedTuple(), NamedTuple(), DefaultContext()), nothing, UnlinkAll(), DynamicPPL.getlogjoint_internal, VarNamedTuple(x = RangeAndTransform{Unlink}(1:1, Unlink()), y = RangeAndTransform{Unlink}(2:2, Unlink())), nothing, 2, [-0.01442938160965795, 1.7571848707614515], DynamicPPL.AccumulatorTuple{3, @NamedTuple{LogPrior::LogPriorAccumulator{Float64}, LogJacobian::LogJacobianAccumulator{Float64}, LogLikelihood::LogLikelihoodAccumulator{Float64}}}((LogPrior = LogPriorAccumulator(0.0), LogJacobian = LogJacobianAccumulator(0.0), LogLikelihood = LogLikelihoodAccumulator(0.0))))

Then you can do:

vals = [1.0, 1.0]
init_strategy = InitFromVector(vals, ldf)

vi = VarInfo()
_, vi = init!!(Xoshiro(468), model, vi, init_strategy, ldf.transform_strategy)
vi

VarInfo
 └─ AccumulatorTuple with 3 accumulators
    ├─ LogPrior => LogPriorAccumulator(-2.3378770664093453)
    ├─ LogJacobian => LogJacobianAccumulator(0.0)
    └─ LogLikelihood => LogLikelihoodAccumulator(-1.4189385332046727)

Partial linking and unlinking

Whole-model link!!(vi, model) and invlink!!(vi, model) remain available. The partial forms, including non-mutating link and invlink, are removed. The second argument of LinkSome and UnlinkSome, fallback, is the strategy for all variables outside vns. The old forms kept their current transforms; pass the strategy that produced vi as fallback to do the same. For example, if y is already linked, LinkSome(Set([@varname(x)]), UnlinkAll()) links x but unlinks y; use LinkSome(Set([@varname(x), @varname(y)]), UnlinkAll()) to keep it linked.

Old:

vi = VarInfo(Xoshiro(468), model)
vns = (@varname(x),)
vi = DynamicPPL.link!!(vi, vns, model)
vi = DynamicPPL.invlink!!(vi, vns, model)

New:

rng = Xoshiro(468)
vi = VarInfo(rng, model)
vns = (@varname(x),)
fallback = UnlinkAll()
linked = LinkSome(Set(vns), fallback)
_, vi = init!!(rng, model, vi, InitFromParams(get_vector_values(vi), nothing), linked)
_, vi = init!!(
    rng,
    model,
    vi,
    InitFromParams(get_vector_values(vi), nothing),
    UnlinkSome(Set(vns), linked),
)
vi

VarInfo
 └─ AccumulatorTuple with 4 accumulators
    ├─ VectorValue => VNTAccumulator{:VectorValue, typeof(DynamicPPL._get_vector_tval), VarNamedTuple{(:x, :y), Tuple{TransformedValue{Vector{Float64}, Unlink}, TransformedValue{Vector{Float64}, Unlink}}}}(DynamicPPL._get_vector_tval, VarNamedTuple(x = TransformedValue{Vector{Float64}, Unlink}([0.07200886749732076], Unlink()), y = TransformedValue{Vector{Float64}, Unlink}([0.9286310592520649], Unlink())))
    ├─ LogPrior => LogPriorAccumulator(-1.769167187291674)
    ├─ LogJacobian => LogJacobianAccumulator(0.0)
    └─ LogLikelihood => LogLikelihoodAccumulator(-0.918938737957795)

To preserve the input vi, pass copy(vi) as the output argument to init!!. Unlike the old partial-link helpers, init!! resets and recomputes all accumulators.

Replacing individual values and transform state

setindex_with_dist!! and update_transform_strategy are removed. Prepare named input values separately and pass the desired strategy to init!!. Use LinkSome, UnlinkSome, or WithTransforms to specify partial or fixed transforms.

Old:

vi = VarInfo(Xoshiro(468), model)
vi = DynamicPPL.setindex_with_dist!!(
    vi, TransformedValue(2.0, NoTransform()), Normal(), @varname(x), nothing
)
_, vi = DynamicPPL.evaluate!!(model, vi)

New:

vi = VarInfo(Xoshiro(468), model)
params = get_vector_values(vi)
params = BangBang.setindex!!(params, TransformedValue(2.0, NoTransform()), @varname(x))
_, vi = init!!(Xoshiro(468), model, vi, InitFromParams(params, nothing), UnlinkAll())
vi

VarInfo
 └─ AccumulatorTuple with 4 accumulators
    ├─ VectorValue => VNTAccumulator{:VectorValue, typeof(DynamicPPL._get_vector_tval), VarNamedTuple{(:x, :y), Tuple{TransformedValue{Vector{Float64}, Unlink}, TransformedValue{Vector{Float64}, Unlink}}}}(DynamicPPL._get_vector_tval, VarNamedTuple(x = TransformedValue{Vector{Float64}, Unlink}([2.0], Unlink()), y = TransformedValue{Vector{Float64}, Unlink}([0.9286310592520649], Unlink())))
    ├─ LogPrior => LogPriorAccumulator(-3.7665745487925504)
    ├─ LogJacobian => LogJacobianAccumulator(0.0)
    └─ LogLikelihood => LogLikelihoodAccumulator(-2.7787474145605433)