[Relay][Training] Add more missing gradients

[altanh]

Added the following gradients:

• take
• reverse_reshape
• stack
• squeeze
• expand_dims
• arange

Also fixed a typo in type solver diagnostics. I had to use a Relay loop in take_grad to support Any size in indices.

cc @t-vi @SWu @MarisaKirisame

[jroesch]

@antinucleon

[antinucleon]

Please try to pass lint locally first~

[MarisaKirisame]

you can get by by defining a 'put' operator, that put a scalar into an index of a tensor, and leave other palces unchanged. put and take has some classic property which I assume will be better for the optimizer. It also allow other optimization (e.g. put and reduce_sum, using grad + (put vala at idxa in 0_array) + (put valb at idxb in 0_array) will be collapsed into a long chain of put on grad, allowing COW to kick in and all take grad mutation update (instead of creating another tensor).

[MarisaKirisame]

@jroesch please look and comment as well.

[altanh]

This is a good point that I was wondering about. The loop is basically just implementing a put operation (like I described in the comment), so it would make sense to have it be a separate op since I imagine it will be useful in general. Should I remove this gradient for now, or keep it and replace it with put once I implement it?

[MarisaKirisame]

Both are fine.

[tqchen]

merging per reviews, feel free to send followup improvements per @MarisaKirisame 's comment

Thanks @jroesch @altanh @antinucleon @MarisaKirisame