Exclude transformed variables from trace

[fonnesbeck]

Skip transformed names when building the expand function output and shape_info.

Following from this, there are changes in the construction of the unconstrained_parameter coordinate. Since shape_info no longer contains transformed names, that broke store_unconstrained=True. Here, _make_functions now also returns (joined_names, joined_shapes) — the full unconstrained parameter space. _prepare_dims_and_coords uses this directly to build the coordinate, decoupling it from shape_info.

[jessegrabowski]

Not a blocker but just to have on your radar is this PR: pymc-devs/pymc#7975

Idea is to be able to transform the labels of dims using the model machinery, so we can label this unconstrained_posterior group. It's a bit of a bummer we lose all the label info currently.

[fonnesbeck]

Idea is to be able to transform the labels of dims using the model machinery, so we can label this unconstrained_posterior group. It's a bit of a bummer we lose all the label info currently.

OK, so the idea is to use transform_labels() on the PyMC side to build proper coordinate labels?

[jessegrabowski]

For now I think we do nothing, I just wanted you to be aware of the issue in case you were interested/motivated to take it over the finish line.

I did a bunch of hack stuff here to get labels on the find_MAP posterior, it's a lot of useless gymnastics but the payoff is nice. I do not suggest using that code, just saying it's an issue we have in a few places now.

[aseyboldt]

That you! Looks good, only one small nitpick: With the normalizing flows, we also have transformed vectors (the parameter space transformed by a flow), so we should not use that word for the constrain/unconstrain distinction in nutpie.

[jessegrabowski]

@fonnesbeck do you plan to have an unconstrained_posterior group in the return, or is that out of scope for this pr

[fonnesbeck]

Not out of scope if we want it. Forgot that we have it in the zarr backend now.

[aseyboldt]

I think that would be great!

[fonnesbeck]

See what you think of this. Tries to mirror PyMC's ZarrTrace convention. I've tried to make this as minimally-invasive as possible, so everything's on the Python side (no Rust). Also dropped the old flat unconstrained_draw vector + unconstrained_parameter coord.

[aseyboldt]

Looks good except for dropping the unconstrained vector. I use this all the time, and would really miss it.
I'm fine with putting it behind a different keyword argument though, the current one is a bit confusing.

[fonnesbeck]

Do you mind them being saved with the same keyword argument? i.e. store_unconstrained stores all unconstrained stuff.

[aseyboldt]

Yes, that sounds good.

[fonnesbeck]

@aseyboldt OK, just uses store_unconstrained= now.

[aseyboldt]

I guess I should have seen this sooner, but can we simplify this by just storing the names of the variables that are unconstrained during compilation, and then just sort them into a new group at the very and of the arviz conversion?

[fonnesbeck]

@aseyboldt OK, unconstrained_posterior group construction now has a simpler design: at compile time we now store only the names of reparameterized value vars; they flow through the trace as ordinary arrow columns and get sorted at the end of _arrow_to_arviz.