custom_grad and vectorized inputs


I’ve been using the NX library, and have been trying to write a custom_grad function that works off of a vectorized input tensor. I’m having trouble getting the function to pass both during autograd phase AND during runtime, usually due to shape issues. I’m not looking for someone to debug my code – instead, I’m wondering if there are any non-trivial examples of custom_grad that I could look at that work off of vectorized tensors.

This one is the example that came to mind for me, but the input is explicitly not vectorized there. Not sure if it helps