How could we (potentially) make use of a (hypothetical) in-place version of Bijectors' API?

[penelopeysm]

In this function

DynamicPPL.jl/src/transformed_values.jl

Lines 392 to 413 in d401df8

	function apply_transform_strategy(
	strategy::AbstractTransformStrategy,
	tv::UntransformedValue,
	vn::VarName,
	dist::Distribution,
	)
	raw_value = get_internal_value(tv)
	target = target_transform(strategy, vn)
	return if target isa DynamicLink
	# Need to link the value. We calculate the logjac
	flink = DynamicPPL.to_linked_vec_transform(dist)
	linked_value, logjac = with_logabsdet_jacobian(flink, raw_value)
	finvlink = DynamicPPL.from_linked_vec_transform(dist)
	linked_tv = LinkedVectorValue(linked_value, finvlink)
	(raw_value, linked_tv, logjac)
	elseif target isa Unlink
	# No need to transform further
	(raw_value, tv, zero(LogProbType))
	else
	error("unknown target transform $target")
	end
	end

We call with_logabsdet_jacobian which allocates a new vector for linked_value. This is all fine and well, because this is pretty much the best we can do with Bijectors' API.

However, let's say that in the long-term future, or maybe with a lot of LLM assistance, we reach a place where an in-place with_logabsdet_jacobian!(f, out, x) function is a guaranteed part of Bijectors' API. (Right now, it isn't.) How can we use that to avoid allocating the intermediate vector?

The main question, IMO, is where does the linked value need to end up in? Right now there are a handful of places where tval is used:

• Inside the values field of a VarInfo. This is hopefully on its way out.
• Inside a VectorValueAccumulator (the direct replacement for the above).
• Inside a VectorParamAccumulator.

The very difficult thing here is that we might want to store the linked_vector value in multiple accumulators. There's no easy way to force them to share memory, not without bringing us back into a VarInfo-like scenario where there is a 'privileged' place to store values in.

To explain what I mean by this, consider that one way of sharing memory between accumulators might be to preallocate any arrays we need, and store those arrays together with the model itself. Then that array would be the central source of truth: with_logabsdet_jacobian! would write into it, and the accumulators would just be views into it. (This would tie in with some ideas currently floating around about having a single, preallocated but empty, VNT stored with the model.) But this is really pretty much just VarInfo reinvented. VarInfo was a single place that collected everything that was going on in the model: for all its imperfections, that was its aim.

There are certain implications that that brings. The main risk of having something that is VarInfo-like is that it risks doing a lot of unnecessary work even when you don't care about it. The beauty of accumulators is that you only need to include what you care about, which means that you don't pay for things you don't need. In contrast, if every model always had a VNT stored with it, and apply_transform_strategy always wrote into it, that would be wasteful. Recall that FastLDF was done pretty much by dispensing with the concept of having a centralised state for model execution.

Basically, I don't have a good answer. In fact, the best answer that I can come up with is that it's very unlikely that this sort of situation will end up being called in a performance-sensitive code path. For the most performance-sensitive bit of DynamicPPL code, namely LogDensityFunction, we don't ever allocate new vectors anyway because no accumulators need them. So it may well be that the answer is that none of this really matters at all, and it's fine for VectorValueAcc or VectorParamAcc to be slightly suboptimal in terms of performance.

But still worth pondering, imo.