pymc.math.stack and xarray dimensions

[hottwaj]

Hi there, I think it would be useful to be able to indicate/annotate that the new dimension along which a set of tensors have been "stacked" by pymc.math.stack is equivalent to a known xarray dimension (previously added by Model.add_coord)

Here's an example - not intended to do anything useful but to illustrate the idea

with pymc.Model() as model:
    model.add_coord('categories', ['A', 'B', C'])
    probs_by_category = pymc.math.stack([var_a, var_b, var_c])

In the above probs_by_category creates a tensor 3 items long. What I would like to do is indicate that probs_by_category has the "categories" xarray dimension attached to it, but as far as I can tell I think it is currently not possible to indicate that. Please correct me if I'm wrong!

Thanks!

[OriolAbril]

probs_by_category as used in the example above is "only" and intermediate result which is not stored by PyMC. Dimensions and coordinates are currently not a feature of tensors themselves but of PyMC variables (and are kept later on in the xarray objects ArviZ uses). Therefore, in order to use dims you need to use pm.Deterministic (which will then also store it as a posterior variable). It would look something like:

with pymc.Model() as model:
    model.add_coord("categories", ["A", "B", "C"])
    probs_by_category = pymc.Deterministic(pymc.math.stack([var_a, var_b, var_c]), dims="categories")

[hottwaj]

Thanks I think you're right that this does work though I'm not working on the model at the moment so can't test. Because xarray dims do not propagate through calculations, there is a bit more boilerplate needed to set dims correctly on every deterministic, and then extra work is needed to change every Deterministic if you change the dimensionality of the model. Even so, getting the right xarray data out of Deterministics like this is very useful, and I understand from other conversations that propagating xarray dims through calculations is not currently feasible (though might be in future if some changes are made on the numpy/xarray side)

Happy to close for now. Thanks!