Add tests for gradient in broadcastable cases

Also check shape of returned arrays in addition to value.

These tests pass on CPU, but do not pass on GPU yet.
I'll investigate tomorrow.