I think I see what you mean by needing to modify each function definition. For now, I’m mostly interested in re-computing the gradient for affine transformations (the addmm(self, mat1, mat2) operation).
Ignoring the alpha and beta parameters for a moment, I noticed that all of the information you need to re-compute the gradient is stored in SavedVariables. Is there a way to access the SavedVariables through the Python API? I ask because the Function object has a saved_variables/saved_tensors 1 property defined in python_function.cpp. However, this property does not seem to exist whenever I try to access it on a Backward object, which is a Function object as I understand it.