A comment describing what `grid_shape` is would be nice.