Added diag! for setting the diagonal of a matrix