Ensure elements in blocks are consecutive

[btalamini]

The element IDs in blocks are not enforced to be consecutive, which is bad for performance. We could sort the elements by block once the mesh is read in or generated. Look into the indices_are_sorted and unique_indices parameters on the jax ndarray indexing tools here.

Another possiblility is to make element block be another global array index, making the loop over blocks be vmap-able. For example, the shape function array would have 3 indices: [block][element][element node].