Fix diag and diagonal.

In the first commit I fixed diag and made minor style changes that other parts of the file use.

The second commit implements a missing dimension check in diagonal that I had forgotten to implement.