Looks great, but doesn't support pytorch yet, only numpy, or does it?

https://github.com/pydata/xarray/issues/3232

my experience was positive in context of pymc and arviz (as a way to represent posterior and posterior predictive distributions). API is sane and quite easy to get around.