You can already use xarray-einstats (and xarray alone in most cases) to compute the loss more easily, but the loss eventually returns a scalar value, so I don’t see how labeled and/or high dimensional data could help on that end.
Maybe we could update the post to pymc v4 and latest ArviZ+xarray. All the iterrows and conversions to datasets are not needed anymore. I think it would make a nice addition to pymc-examples. The loss function in the post is quite simple so there isn’t need for xarray-einstats really, xarray alone is enough.
And if you see examples of more complicated loss functions (i.e. requiring linear algebra or circmean if working with angles) it would also be nice to add a case study about this in xarray-einstats docs.