More rewrites for ExtractDiag

[jessegrabowski]

These came up when I was working on #2032. They're not related to that perse so I split them off. We're missing a lot of simple rewrites for ExtractDiag. I added:

• ExtractDiag(Eye) -> Ones or Zeros, depending on k
• ExtractDiag(Eye * x) -> Alloc(x) (new shape)
• ExtractDiag(Zeros / Ones) -> Zeros / Ones of new shape
• ExtractDiag(Alloc) -> Alloc (with new shape)
• ExtractDiag(Elemwise(a, b)) -> Elemwise(ExtractDiag(a), ExtractDiag(b)) (plus some broadcasting logic)
• ExtractDiag(Transpose(x), offset=k) -> ExtractDiag(X, offset=-k) (I understand transpose is just a free view but it removes a blocker to further rewrites that no longer need to see through the transpose)