Bugfix: diagonal_jacobian after dimension unrolling

Unrolling the dimension loops broke the diagonal_jacobian
transformation since it was looking for index_sums over certain
expressions.