OK, while we're here we should probably fix some problems with set_diagonal:

1. We get a segfault if we try set_diagonal on a matrix where nrows != ncols.
2a. We get a segfault if we try set_diagonal on a mat nest object.
2b. To do so, we need to be able to pull the diagonal blocks out, this is not exposed from petsc4py so that needs some work, and the vec layout will probably be wrong.

1 is straightforward to address and I can do so and add tests.
2a is also straightforward to catch.
2b needs more work, so we should kick it into the long grass.