torch.library.register_fake now accepts more types #124066

zou3519 · 2024-04-15T14:18:51Z

Stack from ghstack (oldest at bottom):

We allow it to accept:

a string with the op name
an opoverload
a new-style custom op

If any of these are referring to a new-style custom op (created with the
custom_op decorator), then we dispatch to CustomOpDef.register_fake.
Otherwise, we do what we previously did.

Test Plan:

new tests

[ghstack-poisoned]

pytorch-bot · 2024-04-15T14:18:55Z

🔗 Helpful Links

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

📄 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 384727b with merge base c4878ab ():
💚 Looks good so far! There are no failures yet. 💚

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

[ghstack-poisoned]

We allow it to accept: - a string with the op name - an opoverload - a new-style custom op If any of these are referring to a new-style custom op (created with the custom_op decorator), then we dispatch to CustomOpDef.register_fake. Otherwise, we do what we previously did. Test Plan: - new tests [ghstack-poisoned]

torch/_library/custom_ops.py

albanD · 2024-04-16T20:20:59Z

torch/_library/custom_ops.py

        yield from check(kwarg)
+
+
+def _maybe_get_opdef(qualname):


qualname sounds like it should be a str object which is not true given the below...

Also it's used only in one place, should we just move it to that other file?

qualname sounds like it should be a str object which is not true given the below...

Ya

Also it's used only in one place, should we just move it to that other file?

I'm trying to not pollute library.py too much

We allow it to accept: - a string with the op name - an opoverload - a new-style custom op If any of these are referring to a new-style custom op (created with the custom_op decorator), then we dispatch to CustomOpDef.register_fake. Otherwise, we do what we previously did. Test Plan: - new tests [ghstack-poisoned]

albanD

Thanks for the update!

We allow it to accept: - a string with the op name - an opoverload - a new-style custom op If any of these are referring to a new-style custom op (created with the custom_op decorator), then we dispatch to CustomOpDef.register_fake. Otherwise, we do what we previously did. Test Plan: - new tests [ghstack-poisoned]

Allows registering autograd for all custom op entry points: - the new-style custom op API (custom_op) - the old-style torch.library APIs - C++ operator registration Test Plan: - tests Pull Request resolved: #124071 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065, #124066

A kernel has "dispatcher convention" if there is an additional keyset arg at the beginning of the argument list. This PR: - adds a way to register kernels with dispatcher_convention using Library.impl (pass dispatcher_convention = True) - adds OpOverload.redispatch We use both of the above in the new custom ops API: we register the autograd kernel in dispatcher convention so that we can actually call redispatch like how pytorch built-in ops do it. Test Plan: - existing tests Pull Request resolved: #124089 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065, #124066, #124071

I forgot to delete this in an earlier PR. Pull Request resolved: #124092 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065, #124066, #124071, #124089

We allow it to accept: - a string with the op name - an opoverload - a new-style custom op If any of these are referring to a new-style custom op (created with the custom_op decorator), then we dispatch to CustomOpDef.register_fake. Otherwise, we do what we previously did. Test Plan: - new tests Pull Request resolved: pytorch#124066 Approved by: https://github.com/albanD ghstack dependencies: pytorch#123937, pytorch#124064, pytorch#124065

Allows registering autograd for all custom op entry points: - the new-style custom op API (custom_op) - the old-style torch.library APIs - C++ operator registration Test Plan: - tests Pull Request resolved: #124071 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065, #124066

A kernel has "dispatcher convention" if there is an additional keyset arg at the beginning of the argument list. This PR: - adds a way to register kernels with dispatcher_convention using Library.impl (pass dispatcher_convention = True) - adds OpOverload.redispatch We use both of the above in the new custom ops API: we register the autograd kernel in dispatcher convention so that we can actually call redispatch like how pytorch built-in ops do it. Test Plan: - existing tests Pull Request resolved: pytorch#124089 Approved by: https://github.com/albanD ghstack dependencies: pytorch#123937, pytorch#124064, pytorch#124065, pytorch#124066, pytorch#124071

I forgot to delete this in an earlier PR. Pull Request resolved: pytorch#124092 Approved by: https://github.com/albanD ghstack dependencies: pytorch#123937, pytorch#124064, pytorch#124065, pytorch#124066, pytorch#124071, pytorch#124089

We allow it to accept: - a string with the op name - an opoverload - a new-style custom op If any of these are referring to a new-style custom op (created with the custom_op decorator), then we dispatch to CustomOpDef.register_fake. Otherwise, we do what we previously did. Test Plan: - new tests Pull Request resolved: #124066 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065

Allows registering autograd for all custom op entry points: - the new-style custom op API (custom_op) - the old-style torch.library APIs - C++ operator registration Test Plan: - tests Pull Request resolved: #124071 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065, #124066

A kernel has "dispatcher convention" if there is an additional keyset arg at the beginning of the argument list. This PR: - adds a way to register kernels with dispatcher_convention using Library.impl (pass dispatcher_convention = True) - adds OpOverload.redispatch We use both of the above in the new custom ops API: we register the autograd kernel in dispatcher convention so that we can actually call redispatch like how pytorch built-in ops do it. Test Plan: - existing tests Pull Request resolved: pytorch#124089 Approved by: https://github.com/albanD ghstack dependencies: pytorch#123937, pytorch#124064, pytorch#124065, pytorch#124066, pytorch#124071

I forgot to delete this in an earlier PR. Pull Request resolved: #124092 Approved by: https://github.com/albanD ghstack dependencies: #123937, #124064, #124065, #124066, #124071, #124089

[WIP] torch.library.register_fake now accepts more

6aaeb88

[ghstack-poisoned]

This was referenced Apr 12, 2024

[custom_ops] Support TensorList inputs/outputs #123615

Closed

Excise old custom ops prototype from custom_op_db #124062

Closed

Update on "torch.library.register_fake now accepts more"

23f1fba

[ghstack-poisoned]

zou3519 changed the title ~~[WIP] torch.library.register_fake now accepts more~~ torch.library.register_fake now accepts more Apr 15, 2024

zou3519 added the suppress-bc-linter Suppresses the failures of API backward-compatibility linter (Lint/bc_linter) label Apr 15, 2024

zou3519 mentioned this pull request Apr 15, 2024

Add torch.library.register_autograd #124071

Closed

zou3519 changed the title ~~torch.library.register_fake now accepts more~~ torch.library.register_fake now accepts more types Apr 15, 2024

This was referenced Apr 15, 2024

Add OpOverload.redispatch; use it in new custom ops API #124089

Closed

Delete qualname from custom_op decorator #124092

Closed

zou3519 mentioned this pull request Apr 16, 2024

Excise uses of the old custom ops APIs #124134

Closed

zou3519 mentioned this pull request Apr 16, 2024

[custom_op] add ability to provide manual schema #124180

Closed

zou3519 added 3 commits April 16, 2024 09:55

This was referenced Apr 16, 2024

Move schema inference to torch._library #124199

Closed

[custom_op] Rename register_impl to register_kernel #124200

Closed

zou3519 requested a review from albanD April 16, 2024 20:02

albanD reviewed Apr 16, 2024

View reviewed changes

zou3519 mentioned this pull request Apr 17, 2024

Add torch.library.register_kernel #124299

Closed

albanD approved these changes Apr 17, 2024

View reviewed changes

zou3519 added ciflow/trunk Trigger trunk jobs on your pull request release notes: composability release notes category labels Apr 17, 2024

pytorchmergebot closed this in 8135c4b Apr 18, 2024

pytorchmergebot added the Merged label Apr 18, 2024

This was referenced Apr 18, 2024

Change register_autograd to reflect ordering of setup_context and backward #124403

Closed

[custom_op] set some tags when constructing the op #124414

Closed

github-actions bot deleted the gh/zou3519/960/head branch June 1, 2024 02:00

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

torch.library.register_fake now accepts more types #124066

torch.library.register_fake now accepts more types #124066

Uh oh!

zou3519 commented Apr 15, 2024 •

edited

Loading

Uh oh!

pytorch-bot bot commented Apr 15, 2024 •

edited

Loading

Uh oh!

Uh oh!

albanD Apr 16, 2024

Uh oh!

albanD Apr 16, 2024

Uh oh!

zou3519 Apr 17, 2024

Uh oh!

albanD left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

torch.library.register_fake now accepts more types #124066

torch.library.register_fake now accepts more types #124066

Uh oh!

Conversation

zou3519 commented Apr 15, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Apr 15, 2024 • 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/124066

✅ No Failures

Uh oh!

Uh oh!

albanD Apr 16, 2024

Choose a reason for hiding this comment

Uh oh!

albanD Apr 16, 2024

Choose a reason for hiding this comment

Uh oh!

zou3519 Apr 17, 2024

Choose a reason for hiding this comment

Uh oh!

albanD left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

zou3519 commented Apr 15, 2024 •

edited

Loading

pytorch-bot bot commented Apr 15, 2024 •

edited

Loading