Hi Richard,
Thanks for your response. Yes I figured this would be the implementation choice, but when I do try to use it, I am told:

RuntimeError: grad can be implicitly created only for scalar outputs

Whereas the version in tf works fine. Do you know what this error might be suggesting?
The shape of the (outputs, wrt inputs)  is ([512,1], [512,3]). They are both variables. Is this the issue?