The reason 2D grids work well is because you can split tensor axes along whatever dim you want across PEs (batch, sequence, hidden, etc.). And your per PE code is relatively simple. It's of the type : get data from neighbor, compute my subtensor, send to neighbor. And you can implement broadcasts and reductions along any axis.