How can I output a gradient in vector format in Op.grad instance?

Hi - I’m encountering a similar problem in generalizing an Op to take vector arguments instead of scalars. I was wondering if you figured out what the problem was in the end. Thanks!