Unensembling

[JeppeKlitgaard]

I am enjoying Equinox a great deal and it seems a huge ergonomic improvement over raw JAX in many cases. Thank you for building it!

I am using Equinox to do ensembling in a similar way to docs#ensambling, but I would like to recover an unensembled version at the end. Similar to the example in the documentation, only some of my properties have gotten batch axes while others were 'broadcast'.

Is there an easy way to split the ensemble into a list of models? My PyTree-foo isn't up to scratch.

I realise this is probably not computationally advantageous in most cases, but for the final part of my processing I would like to leave JAX-land and and recover individual models that I can plot and play with.