Can stop_gradient be factored out of functions?

[cool-RR]

My code is using a complicated loss function which uses stop_gradient on complex expressions. I'm interested in "factoring out" these stop_gradient calls, i.e. rewriting the function to remove those stop_gradient calls while keeping the exact same return values and gradients. I'm thinking of doing that by going over all arguments that affect the expression that is wrapped in stop_gradient, duplicating them with a different name in the function signature, making sure only the duplicate is used in the expression that gets wrapped in stop_gradient, and then removing the stop_gradient calls. Then when calling the function, I'd pass each value twice, once into each of the two arguments.

Do you think this would work?

[mattjj]

Yes, that should work! I think somewhere we wrote down a definition of stop_gradient's semantics in those terms.

[cool-RR]

I can't believe I got it right! I will try this with my code.