Things like linear and relu are at the level of nn. The linear function in particular is composed of subfunctions that may save their input/output as well. Do you explicitly want to only track p q r?
Otherwise, the SavedVariable seem to be the right place to hook in. But it is going to be a big hack and quite fragile