[dynamo][not ready] polyfill infra for classes #146678
Conversation
[ghstack-poisoned]
This PR needs a
|
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
| @substitute_class(functools.partial, supports_reconstruction=True) | ||
| class partial: |
There was a problem hiding this comment.
Can't we use closure to implement class-like objects?
@substitute_in_graph(functools.partial, is_embedded_type=True)
def partial(func, /, *args, **keywords):
def newfunc(*fargs, **fkeywords):
newkeywords = {**keywords, **fkeywords}
return func(*args, *fargs, **newkeywords)
newfunc.func = func
newfunc.args = args
newfunc.keywords = keywords
return newfuncThere was a problem hiding this comment.
In the case of partial, it could be ok. But, I am thinking in terms of more complicated C classes that have many methods. We can still use closure and add method names as attributes to the returned function object. But, I think its more intuitive for the user to have one-to-one mapping for their C++ class to Python class.
Another discussion point (which is not a class vs closure discussion) is how to reconstruct the value. For example, in your above example, if we want to reconstruct the partial object, it will be of type function (and not functools.partial). This causes silent correctness issues because the following code could depend on the returned obj being of type functools.partial. This is not a class vs closure issue though.
There was a problem hiding this comment.
The C++ pytree functions can return a C++ object PyTreeSpec. The pytree polyfills return a Python object with analog methods while setting can_constant_fold_through=False.
See also:
pytorch/torch/_dynamo/polyfills/pytree.py
Lines 119 to 121 in 91c4bf3
pytorch/torch/_dynamo/polyfills/pytree.py
Lines 322 to 334 in 91c4bf3
There was a problem hiding this comment.
For functools.partial:
class partial: # copy the pure-Python version from stdlib
...
@substitute_in_graph(
functools.partial.__new__,
# We need to disable constant folding here because we want the function to reference the
# `partial` class defined above, not the one in the C++ module.
can_constant_fold_through=False,
)
def partial_new(*args, **kwargs):
return partial(*args, **kwargs)
def is_partial_instance(obj):
return instance(obj, partial) # the `partial` class defined aboveOne drawback for polyfills is the isinstance(...) and issubclass(...) checks for polyfilled instances. Neither function-polyfill-with-closure or class-polyfill can resolve this issue elegantly. We also need to provide polyfilled is_partial_instance functions.
There was a problem hiding this comment.
These are really good pointers. I think it makes sense to just polyfill functools.partial.__new__ and provide a helper to convert the already created functools.partial objects to polyfill'd partial objects.
Let me think more about the reconstruction part.
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx chenyang78 kadeng chauhang amjames [ghstack-poisoned]
|
Looks like this PR hasn't been updated in a while so we're going to go ahead and mark this as |
Stack from ghstack (oldest at bottom):
cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @chauhang @amjames