# \[RFC\] New Python operator registration API **URL:** \ **Category:** Uncategorized\ **Created:** [January 25, 2024, 2:58pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838 "2024-01-25T14:58:13Z")\ **Posts on this page:** 11\ **Page:** 1
**Author:** ![zou3519](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/zou3519/32/163_2.png) [@zou3519](https://dev-discuss.pytorch.org/u/zou3519)\ **Post date:** [January 25, 2024, 2:58pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/1 "2024-01-25T14:58:13Z")
We’re developing a new high-level Python operator registration API, designed to be positioned above the existing low-level Python torch.library and C++ TORCH\_LIBRARY APIs. We expect this API to be the first API that users reach for in the future when bringing a custom operator to PyTorch. After some initial prototypes and feedback, we’ve settled on the following design. Please let us know your thoughts. Design doc: [[PUBLIC] Python Custom Ops (2024) - Google Docs](https://docs.google.com/document/d/13CS7nOEtTEsJ_0tiVGPhFDeYRbVb7LXIhw8atS-0eow/edit) ---
**Author:** ![qihqi](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/qihqi/32/284_2.png) [@qihqi](https://dev-discuss.pytorch.org/u/qihqi)\ **Post date:** [January 26, 2024, 6:03pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/2 "2024-01-26T18:03:49Z")
Hi, The goal seems to give another mechanism for defining new Ops, and not to override behavior of existing Aten ops. Is that the intent? Follow up: should potential new core ATen ops (say, scan, while\_loop etc) be implemented this way? ---
**Author:** ![vadimkantorov](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/vadimkantorov/32/1282_2.png) [@vadimkantorov](https://dev-discuss.pytorch.org/u/vadimkantorov)\ **Post date:** [January 26, 2024, 6:35pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/3 "2024-01-26T18:35:24Z")
One my feedback is that it would be nice for such framework to be extensible for some future external op defines corresponding to existing ops (and not only be confined to forward/backward/double-backward). E.g. to be able to define in the future `right_inverse` / `inverse` functions for existing ops and maybe having a higher-order generic functions like : `torch.func.inverse(op)` which would return a function computing an inverse ---
**Author:** ![gilfree](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/gilfree/32/786_2.png) [@gilfree](https://dev-discuss.pytorch.org/u/gilfree)\ **Post date:** [January 30, 2024, 11:25am UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/4 "2024-01-30T11:25:29Z")
Hi, The general approach sounds great, and is simpler then `torch.library` 🙂 There are three things I am missing in this API: 1. (small) Automatic schema deduction from type annotation - In many cases its easy and reduces the effort, and its not hard to verify it against the string like schema. 2. (Maybe out of scope?) Having an option for “native” autograd will be very useful. In my use case I have operators that dynamo cannot handle and cause graph breaks. I want torch.compile to generate a full graph for the forward pass (To have accelerated inference using a custom compiler) but I do not care about the backward speed, and willing to fall back to eager mode for it. I really don’t want to manually write the backward for the operators for these operators. Having something like: 3. (Not sure) Allowing to pass “schema” to the operator. Many operators I use have complex schemas (e.g. quantization schema with multiple config options is selected per layer in the network to give best accuracy). I want to have a way to pass some configuration object. The two options I see now is to serialize to string and deserialize from string (performance? not sure), or to dynamically create a custom op class per configuration (tons of custom ops - almost one per layer) ---
**Author:** ![zou3519](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/zou3519/32/163_2.png) [@zou3519](https://dev-discuss.pytorch.org/u/zou3519)\ **Post date:** [January 30, 2024, 2:13pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/5 "2024-01-30T14:13:51Z")
Yes, the goal is to give a mechanism for defining new ops, not for overriding existing ATen ops. New core ATen ops should still be implemented the way they are implemented inside PyTorch (i.e. in C++ or via the HOP mechanism). ---
**Author:** ![zou3519](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/zou3519/32/163_2.png) [@zou3519](https://dev-discuss.pytorch.org/u/zou3519)\ **Post date:** [January 30, 2024, 2:14pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/6 "2024-01-30T14:14:05Z")
We plan to add more overridable staticmethods for more things that appear in the future. If we need an inverse function, then we’ll add an inverse() method to the class. ---
**Author:** ![zou3519](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/zou3519/32/163_2.png) [@zou3519](https://dev-discuss.pytorch.org/u/zou3519)\ **Post date:** [January 30, 2024, 2:17pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/7 "2024-01-30T14:17:41Z")
@gilfree 1. We’ve thought about this. There are two main problems: (a) which static method do we deduce the schema from? A user can define cpu() for CPU implementatin, cuda() for CUDA implementation, etc. (b) consistency. The schema is not just about type information, it is also about mutability information (e.g. if the operator mutates an input, this shows up in the schema). For consistency, it’s better if there is one way to define the schema (via string) 2. This is an interesting request… but the better solution is probably to use torch.cond in this situation. Creating a custom op from Python means that the code will not work with AOTInductor, so it depends on what your inference pipeline looks like. Using torch.cond to represent both true/false branches means that the code will work in all situations (inference vs training, AOTInductor vs export) 3. What does the configuration object look like? ---
**Author:** ![zou3519](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/zou3519/32/163_2.png) [@zou3519](https://dev-discuss.pytorch.org/u/zou3519)\ **Post date:** [January 30, 2024, 2:18pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/8 "2024-01-30T14:18:50Z")
By the way, if these threads get too long, please feel free to comment directly on the gdoc! Your feedback is very appreciated. ---
**Author:** ![gilfree](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/gilfree/32/786_2.png) [@gilfree](https://dev-discuss.pytorch.org/u/gilfree)\ **Post date:** [January 30, 2024, 3:48pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/9 "2024-01-30T15:48:19Z")
Hi, @zou3519, Thanks for your response. 1 - Ok, got it. I guess that the `abstract` can be used, or even a function named `signature` or something, and add some type annotations e.g. `x: Out[Tensor]`, but it becomes less useful. 2 - I am aware that AOTInducter will not work - I do not plan to use inductor as a backend. We have our own backend - we want the dynamo graph for the forward, in order to compile it with our backend. As for the backward pass - I’m ok with eager - it’s the best I can expect. Dynamo utterly fails on our code - I get ~80 graph breaks per custom convolution layer. Data dependent conditions are just an illustrative example. 3 - We usually work with simple dataclasses, but dicts can also be ok. Thanks Again, If there is anything I can help with pushing this API, let me know, I will do my best. ---
**Author:** ![zou3519](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/zou3519/32/163_2.png) [@zou3519](https://dev-discuss.pytorch.org/u/zou3519)\ **Post date:** [January 30, 2024, 5:20pm UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/10 "2024-01-30T17:20:35Z")
@gilfree 1. It should be simple to build this API on top of the Python Custom Op proposal, but I’m not sure that we’ll build it into the API. What you’d want to do here is just switch on if the input Tensors have requires\_grad=True: if they do, then call a custom op, if they don’t, then call a python function. 2. Is it possible to splat the dataclass so that it can be passed into an operator? ---
**Author:** ![gilfree](https://yyz2.discourse-cdn.com/flex036/user_avatar/dev-discuss.pytorch.org/gilfree/32/786_2.png) [@gilfree](https://dev-discuss.pytorch.org/u/gilfree)\ **Post date:** [January 31, 2024, 8:22am UTC](https://dev-discuss.pytorch.org/t/rfc-new-python-operator-registration-api/1838/11 "2024-01-31T08:22:48Z")
Hi @zou3519 1. Hmm. I think I did bad work explaining myself, will retry: The basic Idea is to have a multi-level compilation. First create a graph at some level of abstraction, do some transformations on it (e.g. conv-bn fuse, but with custom conv) and then pass it to further compilation, or stop at this level and train eagerly with autograd. 2. If by splatting you mean decompose to trivial arguments - not sure. We have some lists there and a large number of members, which will make the schemas complex. Serializing to str is always an option, but in would be much nicer to be able to pass generic python object. Thanks!