Support pointwise marginal log-densities

I propose the factorize keyword of pointwise_logdensities, pointwise_prior_logdensities and pointwise_loglikelihoods be generalized to accept 3 values: nothing, :conditional, :marginal with true => :conditional deprecated but left for backwards-compatibility and same for false => nothing.

Internally, factorize=:conditional continues to call PartitionedDistributions.pointwise_conditional_logpdfs, while factorize=:marginal calls the new PartitionedDistributions.pointwise_marginal_logpdfs.

Background and motivating use-case

For model selection/comparison of a model $M$ via cross-validation with $n$ observations $y$, we can partition the observations into training sets $T_k$ and a non-overlapping evaluation set $E_k$. We can then define a pointwise ELPD estimate as (e.g. https://arxiv.org/abs/2301.08276)

\widehat{\text{elpd}}_{CV}(M | y) = \sum_{k=1}^K \sum_{i \in E_k} w_{ik} \log p(y_i \mid y_{T_k}, M)

for some data-independent weights $w$ (e.g. $w_{ik} = 1$ if the eval sets don't overlap).

Other ELPDs could be written depending on the predictive task. This one though is most comparable to the ELPD estimated by LOO. If e.g. PSIS-LOO fails for many more than 10 observations, it's often going to be less expensive to run 10-fold CV for model selection with this pointwise utility than to do exact LOO for each of the failed observations.

If we get $S$ posterior draws $\theta_s \sim p(\theta \mid y_{T_k})$, then we can estimate the summand with $\text{logsumexp}_{s=1}^S\left(\log p(y_i \mid y_{T_k}, \theta_s) \right) - \log S$. If the observation model is conditionally independent, then currently the quantity inside the log-sum-exp can be computed with pointwise_loglikelihoods with factorize=true.

If the observation model is not factorizable (e.g. an MvNormal with non-diagonal covariance), then computing the desired quantity will probably need to be done manually outside of the model using PartitionedDistributions, something like:

fold_inds = vcat(eval_inds, train_inds)
y_fold = y[fold_inds]
eval_inds_in_fold = first(eachindex(y_fold), length(eval_inds))
likelihood = ... # dist representing p(y_{E_k}, y_{T_k} | theta)
likelihood_eval = conditional(likelihood, y_fold, eval_inds_in_fold)   # dist representing p(y_{E_k} | y_{T_k}, theta)
ll_pw = pointwise_marginal_logpdfs(likelihood_eval, y[eval_inds])  # log p(y_{i} | y_{T_k}, theta) for each i in E_k

This complexity is as far as I see it unavoidable. Since the training set could be much larger then the evaluation set, the construction of the conditional could be quite expensive, and it needs to be done for every Monte Carlo draw $s$ for each fold. However, there's a special case that simplifies. If the the train/eval sets are chosen so the eval obs are conditionally independent from the training obs (i.e. $p(y_{E_k} | y_{T_k}, \theta_s) = p(y_{E_k} | \theta_s)$), then we can do

likelihood_eval = ... # dist representing p(y_{E_k} | theta)
ll_pw = pointwise_marginal_logpdfs(likelihood_eval, y[eval_inds])

It's substantially cheaper but still requires the user to duplicate their likelihood outside of their model, which is error-prone. Or, with the proposal in this issue, they could reuse their model:

model_eval = ... # build model with just `y[eval_inds]`
ll_pw = pointwise_loglikelihoods(model_eval, chains; factorize=:marginal)

Start by tracing pointwise_logdensities, pointwise_prior_logdensities, and pointwise_loglikelihoods and their current factorize handling, then inspect the PartitionedDistributions pointwise conditional and marginal log-density APIs. Done means supporting nothing, :conditional, and :marginal, preserving the true and false compatibility behavior, and using marginal pointwise log densities for the new mode.