Support different dimension permutations

[sethaxen]

The functions currently assume draws are in a single array with shape (draws, chains, params), like MCMCChains.Chains stores. We should consider how to support different permutations (could be as simple as recommending users use PermuteDimsArrays). e.g. ArviZ defaults to (chains, draws, params). Not certain about Soss's SampleChains.

[devmotion]

The functions currently assume draws are in a single array with shape (draws, chains, params)

They either work with vectors of draws or arrays of shape (draws, params, chains).

[sethaxen]

How would one pass a vector of draws to e.g. ess_rhat?

[devmotion]

No, it's not supported by ess_rhat. All functions either work with vectors or arrays of shape (draws, chains, params), but not both.

[ParadaCarleton]

The functions currently assume draws are in a single array with shape (draws, chains, params), like MCMCChains.Chains stores. We should consider how to support different permutations (could be as simple as recommending users use PermuteDimsArrays). e.g. ArviZ defaults to (chains, draws, params). Not certain about Soss's SampleChains.

And I think Stan defaults to (params, chains, draws), giving us a head start on our apparent task of going through every possible permutation of indices.

[devmotion]

I don't think this package should support different permutations. IMO we need a clearly documented and consistent choice for such situations but then users have to permute the arrays if their data is in a different format.

And just to reiterate, not all statistics work with 3d arrays of samples. Some work just with a vector of scalar-valued samples (one parameter, one chain) and I don't think this should be changed. Also the Rstar statistic works with a matrix of samples of size (draws, params) and a corresponding vector of chain indices (this is more general than a 3d array).

[devmotion]

I should add that I don't think we have to stick with the current convention but it would be good to be consistent for 3d arrays. I guess the choice should be motivated by what is the most convenient and efficient layout. Since Julia uses column-major order and Python row-major, probably it differs from what one would choose in Python.

[ParadaCarleton]

I should add that I don't think we have to stick with the current convention but it would be good to be consistent for 3d arrays. I guess the choice should be motivated by what is the most convenient and efficient layout. Since Julia uses column-major order and Python row-major, probably it differs from what one would choose in Python.

I mean, being able to work with different permutations of indices comes as a free side effect of using something like AxisKeys.jl, which I think we should be using anyways to avoid bugs from messing up the order we're indexing in. We can always provide a free "Fallback" that assumes dimensions are ordered in some clearly specified way.

[devmotion]

Please no, let's just stick with generic AbstractArrays - one main motivation here is to get rid of the Chains/AxisArrays mess and just be as generic as possible.

[ParadaCarleton]

I think the best arrangement would be (parameters, draws, chains). Operations on chains are usually done in parallel, e.g. sampling a different chain on each core, so it's not necessary to have chains located close to each other in memory. Parameters are sampled together, so they should be located close in memory so that every time one parameter is written to memory, the next parameter can be written to memory pretty easily. Iterations should be located somewhat close to each other, but aren't always accessed together the same way that parameters from a single iteration usually are. Please correct me if I'm wrong, I'm not a computer scientist.

[ParadaCarleton]

No, it's not supported by ess_rhat. All functions either work with vectors or arrays of shape (draws, chains, params), but not both.

I believe ess_rhat currently uses (draws, params, chains). (An arrangement I find extremely counterintuitive, since chains and draws aren't together.)

[devmotion]

No, it's not supported by ess_rhat. All functions either work with vectors or arrays of shape (draws, chains, params), but not both.

I believe ess_rhat currently uses (draws, params, chains). (An arrangement I find extremely counterintuitive, since chains and draws aren't together.)

Ah yes, this is what I wanted to write and what is mentioned in the documentation. The reason for this layout is that it is the one used in MCMCChains.Chains. However, I don't know what motivated the design choice there, this was decided before I got involved in MCMCChains. Maybe @cpfiffer knows why it was preferred over (parameters, draws, chains)?

[cpfiffer]

I think this was just a holdover from Mamba -- it was really something I hadn't considered.

[sethaxen]

There are two different ways to think about this. One is to reason about what permutation users are likely to pass (which leads to reasoning about what a sensible ordering in memory would be). However, a given PPL may not deliver that ordering (e.g. MCMCChains and IIRC SampleChains). The other is to think about the permutation that would be most efficient for a given function. However, different functions might prefer different permutations, and we should be consistent.

Ideally these two would converge, but I don't know if they do.

[ParadaCarleton]

I think any differences in terms of actual speed/efficiency are probably pretty small -- writing MCMC samples to memory is not going to be the bottleneck for something like HMC under any reasonable set of circumstances. Because of that, I think choices of index orderings should probably be based on what users are most likely to consider intuitive/reasonable orderings. In that case, I think either of (params, draws, chains) or (chains, draws, params) are the most intuitive, since they involve going from more general to more specific, or more specific to more general. (Multiple parameters are contained in a single draw, and several draws are contained in a single chain.) The former has the advantage that it's easier to leave off indices for chains when users are only sampling from a single chain, so I propose we go with that unless anyone has any strong objections.

[ParadaCarleton]

Does anybody object, or should I make a pull request reordering axes like this? I've written some PSIS code that assumes this ordering for ess_rhat, and would like to know whether I should rearrange the PSIS code or the ordering for ess_rhat.