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

username: vdantu
cooked: <p>Hi all,<br>
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.</p>
<p>Also, I wanted to understand why PyTorch JIT IR doesn’t support saving operator node’s attributes?</p>
<p>P.S: This topic is similar to <a href="https://discuss.pytorch.org/t/saving-fusiongroup-as-part-of-scriptmodule/62263" class="inline-onebox">Saving FusionGroup as part of ScriptModule</a></p>

username: Yanan_Cao
cooked: <p>Hi Vamshi,</p>
<p>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.</p>
<p>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.</p>

username: vdantu
cooked: <p>Hi <a class="mention" href="/u/yanan_cao">@Yanan_Cao</a>: Thanks for the response. I had a couple of questions about this</p>
<ol>
<li>I couldn’t find any example of any operator with schema <code>List[Any]</code> in 1.4 or 1.5. Could you please point to any example that I can follow?</li>
<li>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?</li>
</ol>

username: Yanan_Cao
cooked: <ol>
<li>Does it have to be 1.4 or 1.5?</li>
</ol>
<p>Here is an example, not sure if it is in 1.5 or 1.6<br>
</p><aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/pytorch/pytorch/blob/8094228f26cdef24529bffadd6a6e43b56df2d6e/torch/csrc/jit/runtime/register_prim_ops_fulljit.cpp#L674" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/pytorch/pytorch/blob/8094228f26cdef24529bffadd6a6e43b56df2d6e/torch/csrc/jit/runtime/register_prim_ops_fulljit.cpp#L674" target="_blank" rel="nofollow noopener">pytorch/pytorch/blob/8094228f26cdef24529bffadd6a6e43b56df2d6e/torch/csrc/jit/runtime/register_prim_ops_fulljit.cpp#L674</a></h4>
<pre class="onebox"><code class="lang-cpp"><ol class="start lines" start="664" style="counter-reset: li-counter 663 ;">
<li>// round((x_e + x_r)/2)*2 = x_e + round(x_r/2)*2, where x_e is an even integer,</li>
<li>// x_r is either 0.5 of 1.5, round(x_r/2)*2 results a 0 or 2, so the final</li>
<li>// result will always be a even number. Due to symmetricity, it also applies to</li>
<li>// negative cases.</li>
<li>double round_to_even(double a) {</li>
<li>  // NOLINTNEXTLINE(cppcoreguidelines-avoid-magic-numbers)</li>
<li>  return a - std::floor(a) == 0.5 ? (std::round(a * 0.5) * 2.0) : std::round(a);</li>
<li>}</li>
<li>
</li><li>RegisterOperators reg2({</li>
<li class="selected">    // registered as Any[] so that heterogenous tuples can be called with len()</li>
<li>    Operator(</li>
<li>        "aten::len.any(Any[] a) -&gt; int",</li>
<li>        listLen,</li>
<li>        aliasAnalysisFromSchema()),</li>
<li>
</li><li>// these ops have a specialized implementation for the list element type</li>
<li>#define CREATE_SPECIALIZED_LIST_OPS(decl_type, value_type) \</li>
<li>  Operator(                                                \</li>
<li>      "aten::remove." decl_type "(" decl_type              \</li>
<li>      "[](a!) self,                                                           \</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<ol start="2">
<li>I think it can be heterogenous. Indeed it is not a common use case though, let us know if you hit any rough edges.</li>
</ol>

username: vdantu
cooked: <p><a class="mention" href="/u/yanan_cao">@Yanan_Cao</a>: Thanks for the pointers. These were very helpful. Using the type <code>Any</code> seems like a flexible way. It wasn’t supported in 1.4.x though.</p>
<p>I am defining my custom operator as varargs.</p>
<pre><code class="lang-auto">my::Customop(...) -&gt; (...)
</code></pre>
<p>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?</p>

username: Yanan_Cao
cooked: <p>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.</p>