Turns out I can:
```
np.dstack([rdim*rows + i for i in range(rdim)]).flatten()
```