Add an implementation to intercept method calls. #1356
Conversation
A test for extracting gradients with intercept
|
Nice @bastings -- this looks very promising! I skimmed it but haven't yet read the implementation in full detail. A few questions:
|
|
hi!
|
We probably want to keep the separate APIs (at least for now). But I think it's better to keep just one implementation so under the hood I think sow should be implemented as a hook. |
| self.assertEqual(grads['inter_grads']['MyModel_0'] | ||
| ['output_layer']['activation'].shape, (10,)) | ||
|
|
||
| def test_intercept_method_with_multiple_functions(self): |
There was a problem hiding this comment.
please also test intercept method with transformations (vmap and/or scan)
There was a problem hiding this comment.
I have added the vmap test, wasn't sure if you wanted to wrap the apply method with vmap or the whole module with decorator, went with apply. But let me know if that works for you.
| intercept_method: An optional hook that intercepts all function calls and | ||
| can be used to modify args and kwargs before calling each original | ||
| method, and also to modify its output. The provided method needs to | ||
| take the form `f(module, original_function, *args, **kwargs)` and is |
There was a problem hiding this comment.
This signature will fail if users write an intercept_method where kwargs clash with the first 2 arguments. Like (eg.: mdl, mod, fn, fun, etc.) these are all legit names for the first 2 args but could easily happen in kwargs too. An alternative if to make the signature f(module, original_function, args, kwargs) so we can't have any clashes
There was a problem hiding this comment.
Thanks, I went with your suggestion, so the interceptor now takes f(module, original_function, args, kwargs).
Codecov Report
@@ Coverage Diff @@
## master #1356 +/- ##
==========================================
+ Coverage 82.34% 82.37% +0.02%
==========================================
Files 65 65
Lines 5318 5327 +9
==========================================
+ Hits 4379 4388 +9
Misses 939 939
Continue to review full report at Codecov.
|
* Test for intercept with vmap * Comment fix * name fix
jheek
added
Priority: P2 - no schedule
What does this PR do?
Fixes #1355
In particular, it adds an
intercept_methodkwarg toModule.applythat allows users to specify a method that can intercept calls to module methods.Checklist
checks if that's the case).
discussion (please add a
link).
documentation guidelines.
(No quality testing = no merge!)