Hm- what do you think about making comm the first argument (like context on the client-side)?  It's not great that comm is between `shape` and the other distribution parameters like `dist` and `grid_shape`.