How to define and save a custom operator with a dynamic schema?

Hi all,
Is it possible to define and save a custom operator which can potentially take any number of inputs of any type and similarly return any number of outputs of any type? For example, these could be nodes representing/containing subgraph attribute.

Also, I wanted to understand why PyTorch JIT IR doesn’t support saving operator node’s attributes?

P.S: This topic is similar to Saving FusionGroup as part of ScriptModule

Hi Vamshi,

For any number of inputs/outputs of any type, have you tried List[Any]? Your operator would need some way of knowing the exact type of each element though, which can be done by passing some metadata input or attribute.

I am not entirely sure what you mean by “JIT IR doesn’t support saving operator node’s attributes”, since node attributes are saved as far as I know. Please clarify.

Hi @Yanan_Cao: Thanks for the response. I had a couple of questions about this

  1. I couldn’t find any example of any operator with schema List[Any] in 1.4 or 1.5. Could you please point to any example that I can follow?
  2. List by definition expect all elements of same type right? So, how could this be used if the operator takes in “int, bool, float, tensor, lists” in any order and any number?
  1. Does it have to be 1.4 or 1.5?

Here is an example, not sure if it is in 1.5 or 1.6

  1. I think it can be heterogenous. Indeed it is not a common use case though, let us know if you hit any rough edges.

@Yanan_Cao: Thanks for the pointers. These were very helpful. Using the type Any seems like a flexible way. It wasn’t supported in 1.4.x though.

I am defining my custom operator as varargs.

my::Customop(...) -> (...)

This seems to work to save multiple inputs and multiple outputs of different types. Is this a recommended way to represent an operator, or should I look out for any corner case?

1 Like

Given your requirement, I think this is probably the only way to represent the schema of your operator. Yes, you should look out for corner cases given that the version you are using is slightly older.