From dda56ec11810881882461258cb59db15bdfc3130 Mon Sep 17 00:00:00 2001
From: Tor Erlend Fjelde <tor.erlend95@gmail.com>
Date: Fri, 29 Nov 2024 10:00:12 +0100
Subject: [PATCH] use `determine_suitable_varinfo` in `LogDensityFunction` when
 not constructed

---
 src/logdensityfunction.jl | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl
index 9e86590fa..8b6c1769b 100644
--- a/src/logdensityfunction.jl
+++ b/src/logdensityfunction.jl
@@ -65,12 +65,22 @@ end
 
 function LogDensityFunction(
     model::Model,
-    varinfo::AbstractVarInfo=VarInfo(model),
     context::Union{Nothing,AbstractContext}=nothing,
 )
+    # Determine the suitable varinfo for the given model and context.
+    varinfo = determine_suitable_varinfo(model, context === nothing ? leafcontext(model.context) : context)
     return LogDensityFunction(varinfo, model, context)
 end
 
+function LogDensityFunction(
+    model::Model,
+    varinfo::AbstractVarInfo,
+    context::Union{Nothing,AbstractContext}=nothing,
+)
+    return LogDensityFunction(varinfo, model, context)
+end
+
+
 # If a `context` has been specified, we use that. Otherwise we just use the leaf context of `model`.
 function getcontext(f::LogDensityFunction)
     return f.context === nothing ? leafcontext(f.model.context) : f.context