Add a wrapper to transform a NumPy function into a PyTorch function #114610

lezcano · 2023-11-27T18:13:48Z

Stack from ghstack (oldest at bottom):

-> Add a wrapper to transform a NumPy function into a PyTorch function #114610

A less general version of this wrapper was used in the keynote on
torch.compile(numpy). We expose a generic version of the wrapper
that works seamlessly with torch.compile.

cc @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @aakhundov @ColinPeppler

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. [ghstack-poisoned]

pytorch-bot · 2023-11-27T18:13:51Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/114610

📄 Preview Python docs built from this PR
📄 Preview C++ docs built from this PR
❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 79c2403 with merge base a81edf9 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

lezcano · 2023-11-27T18:21:14Z

Tagging a few reviewers as this PR exposes a new API to the public.
This API allows the user to execute NumPy code on CUDA / compute gradients through NumPy code, use NumPy code with torch.func, so I think it is well worth to have it in core and start advertising it.

… function" A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. ghstack-source-id: 105ddf6 Pull Request resolved: #114610

jansel · 2023-11-28T19:23:49Z

docs/source/torch.compiler_faq.rst

+.. code-block:: python
+
+   @torch.compile
+   @torch.compiler.wrap_fn


Should this (and below) be wrap_np? Doesn't seem to match code.

torch/_dynamo/external_utils.py

… function" A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. ghstack-source-id: d0259ae Pull Request resolved: #114610

jansel

Code changes look good to me, though I'd like others to weigh in on if this should be a public API and what it should be called.

albanD · 2023-12-05T22:53:19Z

I'm not sure the semantic of this transformation is very clear to me and when I should use it?

Maybe adding type annotation to the functions will make it easier to see what happens?

lezcano · 2023-12-11T11:29:23Z

The type annotation would be Callable[Any, Any] -> Callable[Any, Any], so I don't think it would add all that much...

albanD · 2023-12-11T16:52:27Z

I mean in terms of Tensor vs array.

… function" A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. ghstack-source-id: 5b49d4d Pull Request resolved: #114610

lezcano · 2023-12-22T09:12:25Z

Sorry, I completely dropped the ball on this one. Added type annotations to all functions. It's not 100% clear to me whether this makes things simpler or not, given that the decorator changes the types, but yeah...

albanD

I think that there is a pretty big constraint here that we end up with a function that can ONLY be used in the context of compile-full-graph and will NOT work in eager.
This is a new enough concept that I think we should say it explicitly in the faq and docstring.

docs/source/torch.compiler_faq.rst

torch/compiler/__init__.py

docs/source/torch.compiler_faq.rst

… function" A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. ghstack-source-id: 43cce18 Pull Request resolved: #114610

lezcano · 2024-01-02T12:24:30Z

Addressed the review. In particular:

I have streamlined the examples in the FAQ, using only wrap_numpy as it's the preferred approach
I have added a paragraph discussing graph breaks explicitly.
I have added a note discussing graph breaks explicitly to the docs.
- Note that the docs are not currently rendered, as the torch.compiler page of the docs is currently being used as a page linking to other pages.

albanD · 2024-01-02T12:47:54Z

docs/source/torch.compiler_faq.rst

 to avoid all these rather expensive memory copies. To do so, we just need
-to tweak our ``numpy_fn`` so that it accepts cuda Tensors and returns tensors:
+to tweak our ``numpy_fn`` so that it accepts cuda Tensors and returns tensors.
+We can do so by using ``torch.compiler.wrap_numpy``:


Thanks for the update, it looks quite good.
Should this piece be split into two questions in the FAQ?

How to speed up NumPy code with torch.compile?

Use the fast cpu backend by enabling torch.compile

Run things on CUDA

How to use NumPy code within your PyTorch code?

The wrap_numpy function

I think splitting it is too much given how little information there is. I will add one example where we simply run things on CPU, implementing your first point "Use the fast cpu backend by enabling torch.compile"

Nevermind, after writing it I chose to split it into two sections, although they are different to the ones you proposed.

… function" A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. cc voznesenskym penguinwu EikanWang jgong5 Guobing-Chen XiaobingSuper zhuhaozhe blzheng wenzhe-nrv jiayisunx peterbell10 ipiszy yf225 chenyang78 kadeng muchulee8 aakhundov ColinPeppler [ghstack-poisoned]

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. ghstack-source-id: 2e17422 Pull Request resolved: #114610

albanD

Thanks for the update!

lezcano · 2024-01-02T16:00:14Z

@pytorchbot merge

pytorchmergebot · 2024-01-02T16:02:08Z

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging

Check the merge workflow status
here

Add a wrapper to transform a NumPy function into a PyTorch function

7108cd8

A less general version of this wrapper was used in the keynote on `torch.compile(numpy)`. We expose a generic version of the wrapper that works seamlessly with `torch.compile`. [ghstack-poisoned]

lezcano mentioned this pull request Nov 27, 2023

Allow to differentiate through NumPy code #114608

Closed

lezcano mentioned this pull request Nov 27, 2023

[easy] Dispatch torch.from_numpy to torch.as_tensor #114609

Closed

github-actions bot added module: inductor module: dynamo ciflow/inductor labels Nov 27, 2023

pytorchbot added the open source label Nov 27, 2023

lezcano requested review from albanD, ezyang, jansel and malfet November 27, 2023 18:18

lezcano added the release notes: dynamo label Nov 27, 2023

lezcano requested a review from soumith November 27, 2023 22:20

jansel requested changes Nov 28, 2023

View reviewed changes

jansel reviewed Nov 29, 2023

View reviewed changes

ezyang removed their request for review January 1, 2024 01:42

albanD reviewed Jan 2, 2024

View reviewed changes

docs/source/torch.compiler_faq.rst Outdated Show resolved Hide resolved

torch/compiler/__init__.py Outdated Show resolved Hide resolved

docs/source/torch.compiler_faq.rst Show resolved Hide resolved

albanD reviewed Jan 2, 2024

View reviewed changes

albanD approved these changes Jan 2, 2024

View reviewed changes

pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Jan 2, 2024

pytorchmergebot added the merging label Jan 2, 2024

pytorchmergebot added Merged and removed merging labels Jan 2, 2024

pytorchmergebot closed this in b18d8d4 Jan 2, 2024

facebook-github-bot deleted the gh/Lezcano/256/head branch January 6, 2024 15:22

Add a wrapper to transform a NumPy function into a PyTorch function #114610

Add a wrapper to transform a NumPy function into a PyTorch function #114610

Uh oh!

Conversation

lezcano commented Nov 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Nov 27, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/114610

✅ No Failures

Uh oh!

lezcano commented Nov 27, 2023

Uh oh!

jansel Nov 28, 2023

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jansel left a comment

Choose a reason for hiding this comment

Uh oh!

albanD commented Dec 5, 2023

Uh oh!

lezcano commented Dec 11, 2023

Uh oh!

albanD commented Dec 11, 2023

Uh oh!

lezcano commented Dec 22, 2023

Uh oh!

albanD left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lezcano commented Jan 2, 2024

Uh oh!

albanD Jan 2, 2024

Choose a reason for hiding this comment

Uh oh!

lezcano Jan 2, 2024

Choose a reason for hiding this comment

Uh oh!

lezcano Jan 2, 2024

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

lezcano commented Jan 2, 2024

Uh oh!

pytorchmergebot commented Jan 2, 2024

Merge started

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

lezcano commented Nov 27, 2023 •

edited

Loading

pytorch-bot bot commented Nov 27, 2023 •

edited

Loading

albanD left a comment •

edited

Loading