`pm.LKJCorr` breaks `compute_log_prior`/`compute_log_likelihood` on a value-variable naming mismatch, blocking psense for `rlm-jc-001`

[frankbuckley]

Note

Drafted by an LLM-based AI tool (Claude Code/Opus 5).

Summary

pm.compute_log_prior and pm.compute_log_likelihood both refuse to run on any model containing a pm.LKJCorr random variable, because they match the trace's posterior variables against the model's transformed value variable names while the posterior (correctly) stores the constrained name. The mismatch is a naming seam, not an intractable density.

Today this costs exactly one fit — lrp-rlm-jc-001 is the only model that both uses LKJCorr and routes through the standard compute_log_* path — so it is the last outstanding item on #381's psense coverage after #452 (192 of 194 fits covered; lrp-rlm-mm-001 is the other, and is a genuine non-converged exemption). But it will bite any future LKJCorr model on the standard path, and #383 recommends moving more models to LKJCorr, so the seam is worth closing before it spreads.

Reproduction

import pymc as pm

with pm.Model() as m:
    L = pm.LKJCorr("measure_corr_chol", n=3, eta=2.0)

m.free_RVs    # [measure_corr_chol]
m.value_vars  # [measure_corr_chol_cholesky_corr__]

LKJCorr's default transform is a Cholesky-corr transform, so the value variable is <name>_cholesky_corr__. The posterior group stores the constrained draw under <name>. compute_log_prior derives the name it expects from the transform and finds no match:

exact match required for all data variable names, but
['eta_cell', 'sigma_subject', 'measure_corr_chol_cholesky', 'z_subject', 'kappa'] !=
['eta_cell', 'sigma_subject', 'kappa', 'measure_corr_chol', 'z_subject']
{'measure_corr_chol_cholesky', 'measure_corr_chol'} are not in both.

Observed on a real --config reporting fit of lrp-rlm-jc-001 on 2026-07-26. Both groups fail — the prior group as well as the likelihood group — so this is not the "likelihood is not evaluable" case that compute_log_likelihood_and_prior(strict=False) was written for. That argument lets the fit survive the failure; it does not repair it.

Why only one model is affected today

pm.LKJCorr appears three times in factories.py:

Line	RV	Model(s)	Affected?
5641	trait_corr_chol	lcf (longitudinal correlated-factor)	No
5646	state_corr_chol_w{n}	lcf	No
6319	measure_corr_chol	rlm-jc-001 (historical joint)	Yes

lcf escapes because it builds its log_likelihood / log_prior groups by a bespoke path rather than calling compute_log_likelihood_and_prior — a design already documented in that function's docstring ("Families with a bespoke log-likelihood … build their groups their own way and do not call this"), and confirmed by lrp-rli-lcf-001's stored trace carrying both groups and a valid psense_summary.csv. fit_rlm_joint_growth is compute_loo=False, so the sampling stage attaches nothing and it must request the groups explicitly — which is where it meets the seam.

Note the choice of LKJCorr here is deliberate and should not be reverted: factories.py:6315-6318 records that LKJCholeskyCov was rejected in this correlation-only role because its unused per-outcome sd scales are unidentified nuisances that pollute the all-free-RV convergence gate. #383 makes the same argument for switching mm / rlm-mm to LKJCorr, which is precisely why this seam should be fixed rather than worked around by going back.

Options

1. Give historical_joint a bespoke log-prior path, as lcf already has. Most tractable, has in-repo precedent, and touches one family. Downside: a second bespoke path to maintain, and it does not help the next LKJCorr model.
2. Reconcile the name at the compute_log_* seam — pass the constrained-to-value mapping (model.rvs_to_values) or rename the trace variables into a throwaway copy before the call. Fixes every present and future LKJCorr model in one place. Needs care that the values passed are on the scale the log-density expects, which is the substantive risk: a silently wrong scale would produce plausible numbers rather than an error.
3. Check whether a newer PyMC resolves it upstream before building anything. Worth ten minutes first — if it is fixed above 6.1, this becomes a version bump plus a regression test.

Option 3 first, then 2 if the seam is ours to close, with 1 as the fallback.

Acceptance criteria

• Root cause confirmed as ours or upstream, with the PyMC version behaviour recorded.
• lrp-rlm-jc-001 emits psense_summary.csv and psense.png at --config reporting, or the model is recorded as a documented exemption with the specific reason (matching the treatment lrp-rlm-mm-001 already gets under Emit power-scaling (psense) and prior-pushforward artifacts for all estimand-reporting families #381), not left silently absent.
• If the fix is at the seam (option 2), a test asserts that a model containing an LKJCorr RV can have compute_log_prior run against its own posterior — the naming contract, pinned so a PyMC upgrade cannot silently reintroduce it.
• If the fix changes which values reach the log-density, the resulting psense numbers are sanity-checked against a model whose psense already works (lrp-rli-lcf-001 is the natural comparator — same LKJCorr construction, working groups).
• notes/202607261700-psense-coverage-backfill.md updated with the outcome; Emit power-scaling (psense) and prior-pushforward artifacts for all estimand-reporting families #381's coverage line moves from 192/194.
• ruff check src/, npm run format:check, npm run spellcheck pass.

Related

Found while closing #381's refit gap in #452, which added the psense wiring to fit_rlm_joint_growth that surfaced this — before that the artefacts were simply absent with no error at all. Connected to #383 (which would move more models onto LKJCorr) and to the general LKJCholeskyCov-vs-LKJCorr reasoning recorded at factories.py:6315-6318 and in notes/202607101638-mm-001-convergence-reparameterisation.md.

[frankbuckley]

Note

Drafted by an LLM-based AI tool (Claude Code/Opus 5).

Upstream check: PyMC 6.2.0 does not fix this, and option 3 is dead

Checked as the first acceptance item. We are on PyMC 6.1.0; 6.2.0 is the latest release. The relevant code is byte-identical in both, so an upgrade changes nothing here.

Exact root cause

pymc/util.py:

def get_transformed_name(name, transform):
    return f"{name}_{transform.name}__"

def get_untransformed_name(name):
    if not is_transformed_name(name):
        raise ValueError(f"{name} does not appear to be a transformed name")
    return "_".join(name.split("_")[:-3])          # <- assumes transform.name has no underscore

get_untransformed_name drops a fixed three trailing underscore-separated components — one for the transform name, two for the __ marker. That is only correct when transform.name itself contains no underscore. LKJCorr's default transform is CholeskyCorrTransform, whose name is "cholesky_corr", so one component of the transform name survives the strip:

transform	value var	round-trips to
log	myvar_log__	myvar	✅
logodds	myvar_logodds__	myvar	✅
ordered	myvar_ordered__	myvar	✅
sumto1	myvar_sumto1__	myvar	✅
zerosum	myvar_zerosum__	myvar	✅
cholesky-cov	myvar_cholesky-cov__	myvar	✅
cholesky-cov-packed	myvar_cholesky-cov-packed__	myvar	✅
cholesky_corr	myvar_cholesky_corr__	myvar_cholesky	❌
log_exp_m1	myvar_log_exp_m1__	myvar_log_exp	❌

Reproduced against the installed 6.1.0 with pymc.util.get_transformed_name / get_untransformed_name directly, and confirmed unchanged at the v6.2.0 tag.

Two observations worth carrying:

• LKJCorr is not the only affected distribution. LogExpM1 (name = "log_exp_m1", two underscores) breaks the same way. We do not currently use it, but the bug is a property of the name parser, not of LKJCorr.
• The hyphenated names look like a prior workaround. cholesky-cov and cholesky-cov-packed use hyphens where cholesky_corr uses an underscore, and hyphens round-trip cleanly. That reads as someone having hit this before and dodged it by naming convention rather than fixing the parser — which is why the seam is still open for the one transform that did not get the treatment.

Where it surfaces

compute_log_density (pymc/stats/log_density.py) builds its elemwise function over remove_value_transforms(model).value_vars but subsets the posterior by [rv.name for rv in model.free_RVs]. remove_value_transforms → change_value_transforms renames each value variable via get_untransformed_name(value.name), so the LKJCorr value var comes out as measure_corr_chol_cholesky while the posterior carries measure_corr_chol. xarray then raises the exact-match error, for the prior group as well as the likelihood group.

Revised options

Option 3 (wait for upstream) is dead. That leaves:

1. Fix at the PyMC seam and upstream it. In change_value_transforms the loop already holds rv, value = node.inputs, so the untransformed name is available directly as rv.name rather than being re-derived by string surgery on value.name. Using it would fix every underscore-named transform at once and remove the parser's implicit naming constraint. This needs checking against PyMC's own tests — get_untransformed_name is public API and other call sites may depend on the string behaviour — but it is the fix that addresses the cause rather than the symptom. If it holds, worth a PR to pymc-devs/pymc.
2. Work around it locally for historical_joint, either by a bespoke log-prior path as lcf already has, or by renaming the offending posterior variable into a throwaway copy before calling compute_log_*. Cheaper and self-contained, but leaves the trap set for the next LKJCorr model.

Recommend 1 with 2 as the interim, since 1 depends on an upstream release we do not control.

Acceptance criteria (updated)

• Root cause confirmed as upstream, with version behaviour recorded: present in 6.1.0 and 6.2.0, caused by get_untransformed_name's fixed three-component strip.
• Decide between the upstream fix and the local workaround (and whether to open a PyMC issue/PR).
• lrp-rlm-jc-001 emits psense_summary.csv and psense.png at --config reporting, or is recorded as a documented exemption with this specific reason.
• A regression test pins the naming contract so a future PyMC upgrade cannot silently reintroduce or mask it.
• If the fix changes which values reach the log-density, sanity-check the resulting psense against lrp-rli-lcf-001 (same LKJCorr construction, working groups).
• notes/202607261700-psense-coverage-backfill.md updated; Emit power-scaling (psense) and prior-pushforward artifacts for all estimand-reporting families #381's coverage line moves from 192/194.
• ruff check src/, npm run format:check, npm run spellcheck pass.