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.