Dispatching non-leaf nodes

username: almeetb
cooked: <p>Hello,</p>
<p>I have a backend that has registered ops in in <code>RegistrationDeclarations.h</code> and is receiving dispatched python operations but for some reason, I do not get non-leaf nodes (<code>compound: true)</code> dispatched. For e.g. the following forward function of my module</p>
<pre><code class="lang-auto">def forward(self, x):
   return torch.nn.functional.selu(x)
</code></pre>
<p>Dispatches elu instead of selu and does not dispatch any mul for the scale factor.<br>
I am using pytorch 1.5</p>

username: albanD
cooked: <p>Hi,</p>
<p>Could you be more precise about what you want to do and how you implement it?<br>
Are you registering new ops? Are you modifying pytorch source directly? How do you do that?</p>

username: almeetb
cooked: <p>I want to do something similar to Pytorch/XLA backend. I modified Pytorch source to recognize a new backend and device type and in my custom c++ extension I registered the op by:</p>
<pre><code class="lang-auto">static auto registerer = torch::RegisterOperators()
   .op(torch::RegisterOperators::options().schema("aten::selu(Tensor self) -&gt; Tensor")
       .impl_unboxedOnlyKernel&lt;at::Tensor(const at::Tensor &amp;), &amp;MyRegisteredOps::selu&gt;(at::DispatchKey::NewDispatchKey)
</code></pre>
<p>I am not creating new ops, just specifying a different backend for current ops but I don’t see any non-leaf nodes being dispatched such as the selu example above.</p>

username: albanD
cooked: <p>Hi,</p>
<p>Well it depends where your dispatch key is wrt to the autograd one. If it is below, then the autograd will be taken care of before getting to your code.</p>
<p>Also note that xla is out of tree and doing something similar (shouldn’t) doesn’t require changing core torch. You can use for example <code>PrivateUse1</code> that is especially design to allow users to have their own backend.</p>

username: almeetb
cooked: <p>thanks I will uplift to using <code>PrivateUse1</code><br>
but could you explain how taking care of the autograd key results in non-leaf nodes not being dispatched? If my dispatch key is above, would I still be able to use pytorch autograd engine?</p>

username: albanD
cooked: <aside class="quote no-group" data-username="almeetb" data-post="5" data-topic="92445">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/34f0e0/48.png" class="avatar"> almeetb:</div>
<blockquote>
<p>non-leaf nodes not being dispatched?</p>
</blockquote>
</aside>
<p>I’m not sure what you mean by that.<br>
Can you give a concrete example?</p>

username: almeetb
cooked: <p>The selu vs. elu example is the one I see manifesting. selu is not a leaf op so even after registering it I only get elu.</p>
<p>If the PrivateUse1 is listed after the autograd dispatch key would that result in selu being dispatched?</p>

username: albanD
cooked: <p>Ho so you’re trying to register new dispatch keys for operators that are already defined in core pytorch?<br>
We use the term “leaf” in the autograd usually but it is unrelated here <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Yes I think that some ops that are just implementing a generic function don’t go through the dispatcher at all. And so you registering new impl for them won’t work.</p>
<p>cc <a class="mention" href="/u/sebastian_messmer">@Sebastian_Messmer</a> who will know better.</p>

username: almeetb
cooked: <p>thanks! I saw that some operations just dispatch the decomposed variants but the elu vs selu led me to believe that in some cases I could be getting the wrong ops. <a class="mention" href="/u/sebastian_messmer">@Sebastian_Messmer</a></p>

username: Sebastian_Messmer
cooked: <p>You should see both calls, first to selu, and then to elu. Our implementation of selu just calls elu (with a scale parameter, that’s why you don’t see a mul). As long as you only do inference you should be able to register a kernel for selu and intercept that first call. If you want to do Autograd, it gets a bit more complicated.</p>

username: almeetb
cooked: <p>How would I intercept the first call and do you mind explaining what would happen in the case I want to do Autograd?<br>
I am also seeing a case where the forward function calls <code>torch.flatten(x)</code> but I don’t see the flatten</p>
<p>Thanks for your help!</p>

username: Sebastian_Messmer
cooked: <p>The standard operator registration should work. Something like:</p>
<pre><code class="lang-auto">static auto registry = torch::RegisterOperators().op("aten::selu", torch::RegisterOperators::options()
   .kernel(DispatchKey::PrivateUse1, &amp;your_selu_implementation)
);
</code></pre>
<p>We’re currently thinking about a redesign that might make autograd for these cases easier in the future, but for now it’s complicated. You could try and follow what XLA is doing. I’ll get back to you.</p>

username: ailzhang
cooked: <p>Hi <a class="mention" href="/u/almeetb">@almeetb</a>, in pytorch dispatcher <code>compound=True</code> means we have a mathematical definition of <code>selu</code> that works for all backends and we just by default use that implementation. So even you register a kernel to <code>selu</code> it won’t get called, instead it’ll be called at <code>elu</code>which is a non-compound op.<br>
In general this works fine for most XLA ops (and it saves a lot of time as compound ops are supported “automatically”). But if you really really want to override a compound kernel, take a look at <a href="https://github.com/pytorch/xla/blob/master/torch_xla/csrc/aten_autograd_ops.cpp#L18" rel="nofollow noopener">https://github.com/pytorch/xla/blob/master/torch_xla/csrc/aten_autograd_ops.cpp#L18</a>. You’ll need to override both forward and backward through torch::autograd::Function, and register that kernel to a <code>XLAPreAutograd</code> key.</p>

username: Sebastian_Messmer
cooked: <p>Ailing is correct. My previous statement about you being able to see the first call to selu is only true if you have autograd disabled. If you have autograd enabled, then even in the forward path, you will not see your custom selu unless you do the workaround proposed by Ailing.</p>
<p>The recommended way for now is to not intercept compound ops and instead override leaf nodes like elu.</p>

username: almeetb
cooked: <p>I see, thank you and <a class="mention" href="/u/ailzhang">@ailzhang</a> for your help!</p>