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

vdantu · August 9, 2020, 7:56pm

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

Yanan_Cao · August 11, 2020, 4:20pm

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.

vdantu · August 11, 2020, 5:04pm

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

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?
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?

Yanan_Cao · August 19, 2020, 6:02am

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

github.com

pytorch/pytorch/blob/8094228f26cdef24529bffadd6a6e43b56df2d6e/torch/csrc/jit/runtime/register_prim_ops_fulljit.cpp#L674


// round((x_e + x_r)/2)*2 = x_e + round(x_r/2)*2, where x_e is an even integer,
// x_r is either 0.5 of 1.5, round(x_r/2)*2 results a 0 or 2, so the final
// result will always be a even number. Due to symmetricity, it also applies to
// negative cases.
double round_to_even(double a) {
  // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)
  return a - std::floor(a) == 0.5 ? (std::round(a * 0.5) * 2.0) : std::round(a);
}

RegisterOperators reg2({
    // registered as Any[] so that heterogenous tuples can be called with len()
    Operator(
        "aten::len.any(Any[] a) -> int",
        listLen,
        aliasAnalysisFromSchema()),

// these ops have a specialized implementation for the list element type
#define CREATE_SPECIALIZED_LIST_OPS(decl_type, value_type) \
  Operator(                                                \
      "aten::remove." decl_type "(" decl_type              \
      "[](a!) self,                                                           \

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

vdantu · August 25, 2020, 11:41pm

@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?

Yanan_Cao · August 25, 2020, 11:59pm

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.