> @rhattersley I've pushed a new commit which avoids the need to signal three states with one variable, and moves the dimensionality check right next to where it's used. Hopefully that does the trick for you.

Great, much prefer this.
I think you could now lose the extra variable 'dim', as it's only used once + I think the logic is quite clear without it.