Hey,
I looks like none of your formulas is using the grad_output parameter?
What the backward function for a functin f should be computing is the vector Jacobian product between grad_output and the Jacobian of f.