Remove LogDensityFunctionWrapper and replace VarInfo(mld, ...) with InitFromVector

[anurag-mds]

LogDensityFunctionWrapper was a trivial struct whose sole purpose was to convert the LogDensityFunction type to the two-argument callable signature expected by MarginalLogDensities, and to carry a VarInfo to be reconstructed later. But LogDensityFunction already contains all the information we'll ever need (it contains both the model and the transform_strategy as well as _varname_ranges), so this wrapper was unnecessary. I removed it and added a two-argument callable method to LogDensityFunction in the extension.

The VarInfo(mld, ...), which this method enabled, was also problematic because it would unflatten parameter values into a VarInfo using unflatten!!, which would leave the log probabilities and other accumulated data in an inconsistent state. I've now replaced it with InitFromVector(mld, ...), which returns an initialisation strategy that the user can pass to init!! to get a fully consistent VarInfo by re-computing the model.

Closes #1308

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 78.58%. Comparing base (8d1de99) to head (98c8af4).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1335      +/-   ##
==========================================
- Coverage   78.59%   78.58%   -0.02%     
==========================================
  Files          50       50              
  Lines        3626     3624       -2     
==========================================
- Hits         2850     2848       -2     
  Misses        776      776              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[yebai]

I’d suggest we keep VarInfo.

[penelopeysm]

I wholly disagree -- the varinfo generated contained wrong data, and it's not something we should be giving people to use.

[yebai]

Ah, okay, maybe I didn't fully get the motivation. Let's chat when we meet.

[anurag-mds]

Actually I used the JuliaFormatter once directly as

using JuliaFormatter
format(".")

This changed multiple files and also throwed many warnings saying The formatter couldn't be applied here.

then I ran for the specific files I changed

using JuliaFormatter
format(["HISTORY.md", "docs/make.jl", "docs/src/api.md", "ext/DynamicPPLMarginalLogDensitiesExt.jl", "test/ext/DynamicPPLMarginalLogDensitiesExt.jl"])

the formatter worked and I pushed the code but still formatting CI is failing I can't get it why?
Can you help me?

[penelopeysm]

Are you using JuliaFormatter v1? https://turinglang.org/docs/contributing/code-formatting/

[yebai]

I am concerned that InitFromVector is not a super clear name and that we may need to rename it later, whereas VarInfor is more stable.

[penelopeysm]

It already exists, though: https://turinglang.org/DynamicPPL.jl/stable/ldf/models/ and it's not really that complicated to use.

I'm conscious of backwards compatibility, but I really strongly believe that VarInfo is not the right thing to keep around. As you are aware, there have been many perf gains in DynamicPPL recently. But we haven't managed to do that because I'm some hacker Linus Torvalds genius who optimised low level code: they're just because we have better, more modular, interfaces that don't make users pay for things that they don't need.

VarInfo is the exact antithesis of that: it's a struct that bundles lots of stuff together and forces users to pay for things they don't need. Besides the performance issues, it also leads to incorrect results with unflatten!! because log-probs and transforms are not updated in tandem with the parameters, which is an issue I've pointed out many times. That's exactly the reason why I think we should remove this method.

If we want Turing/DynamicPPL to be something that people trust, then correctness is very, very important. I think it's the single most important thing, even more so than performance.

[anurag-mds]

Are you using JuliaFormatter v1? https://turinglang.org/docs/contributing/code-formatting/

Fixed: was using the wrong JuliaFormatter version. Formatting CI passed now.

[anurag-mds]

I haven't commented yet as I was following the discussion between the maintainers regarding the interface design.

I just wanted to clarify: was the change mentioned by @yebai intended for me to implement, or was it a request for @penelopeysm’s approval?

If it is for me, I will update the name from InitFromVector to VarInfor Is there any other change I should include?

Could you also clarify if VarInfor was a typo? I assume you meant the standard VarInfo.

[penelopeysm]

It's a typo. And it's not just about the name, it's really about the interface design. I don't really want to repeat myself, so I'll just point to #1308 (comment).

[anurag-mds]

Hi @penelopeysm, I ran into two issues:

Issue 1: Base.Fix1 doesn't work
When using Base.Fix1(LogDensityProblems.logdensity, ldf), I get:

MethodError: no method matching logdensity(::LogDensityFunction, ::Vector{...}, ::Tuple{})

Because MarginalLogDensities calls the callable with two arguments (x, extra_data), so Base.Fix1 expands to logdensity(ldf, x, extra_data) 3 args instead of 2.

My suggestion is to use an anonymous function instead:

julia(x, _) -> LogDensityProblems.logdensity(ldf, x)

This captures ldf without making LogDensityFunction globally callable, and correctly absorbs the second argument. Could you confirm if this is the right approach?

Issue 2: Tracer types -- <:Real constraint
MarginalLogDensities passes tracer/dual number types during Hessian computation, but logdensity only accepts AbstractVector{<:Real} and tracer types are not <:Real. I removed the <:Real constraint in 3 places in src/logdensityfunction.jl (logdensity, logdensity_at, and LogDensityAt's callable), and confirmed from the error logs that the closest candidate correctly changed from AbstractVector{<:Real} to AbstractVector.

Could you confirm whether this is the right fix?

[penelopeysm]

Ahh, yes, darn. We still need some kind of wrapper around LogDensityFunction. You can use an anonymous function, I think I generally prefer a callable struct though because (1) it's easier to dispatch on in type signatures, and (2) also easier to compare (every new anonymous function declaration creates a separate one).

julia> (x -> x + 1) == (x -> x + 1)
false

So the only new thing is that the wrapper doesn't store a VarInfo anymore, sorry to mislead on that.

The tracer types should subtype Real. What are the types that you are having trouble with?

[anurag-mds]

Hey @penelopeysm, Can you check the PR?

[penelopeysm]

I'll look at it next week when I'm back at work, but I'll self-request a review so that it doesn't slip off the radar.