Feature request: Named RV dimensions

[twiecki]

Ekatarina just had a good idea: naming dimensions of random variables:

pm.Normal("mu", group_mu["control_group"], 1) instead of pm.Normal("mu", group_mu[0], 1).

[OriolAbril]

I think this is closely tied to the discussion on #4552 which will provide useful context, especially on the difference between dimensions and coordinate labels. The example above is actually about using labeled indexing with coordinate labels. In positional/numpy-like indexing, the dimension is defined by the position within the [] (i.e. [0, 3] gets the first position of the first dimension and the 4th position of the 2nd dimension whereas [3, 0] gets the 4th position of the 1st dimension and the 1st position of the 2nd). In xarray these would translate to {"dim1": 0, "dim2": 3} and {"dim1": 3, "dim2": 0} (or {"dim2": 0, "dim1": 3}) respectively.

From the example I understand we want to get the position labeled control_group of the 1st (and only) dimension in group_mu.

[michaelosthege]

That's like a smart translation from coord values into indices, right?
So given

coords = {
    "treatment": ["juice", "ethanol", "caffeine"],
}
effect_size = pm.Normal("effect_size", dims=("treatment"))

We could write

effect = pm.Deterministic(effect_size["caffeine"])

Instead of

effect = pm.Deterministic(effect_size[coords["treatment"].index("caffeine")])

I am not familiar with how we treat dims/coords right now, or how that changes with RandomVariables.

[twiecki]

Related: pymc-devs/pytensor#352.

[ricardoV94]

Closing as this is being discussed in the context of PyTensor

[OriolAbril]

Can you link to the related issue in pytensor if there is one?

[ricardoV94]

@aseyboldt did you write down something already?