def reassemble_array(arr):
    """This converts an array containing the solution from each processor as
    first axis to the full solution.

    Shape (mpi_size, ny_local, nx_local) -> (ny_global, nx_global)
    """
    out = jnp.empty((ny_global, nx_global), dtype=arr.dtype)
    for i in range(mpi_size):
        proc_idx_i = np.unravel_index(i, (nproc_y, nproc_x))
        local_slice_i = (
            slice(
                (ny_local - 2) * proc_idx_i[0],
                (ny_local - 2) * proc_idx_i[0] + ny_local,
            ),
            slice(
                (nx_local - 2) * proc_idx_i[1],
                (nx_local - 2) * proc_idx_i[1] + nx_local,
            ),
        )
        out = out.at[local_slice_i].set(arr[i])

    return out
This converts an array containing the solution from each processor as
    first axis to the full solution.

    Shape (mpi_size, ny_local, nx_local) -> (ny_global, nx_global)
    
examples/shallow_water.py