Hello,
I have a backend that has registered ops in in RegistrationDeclarations.h and is receiving dispatched python operations but for some reason, I do not get non-leaf nodes (compound: true) dispatched. For e.g. the following forward function of my module
def forward(self, x):
return torch.nn.functional.selu(x)
Dispatches elu instead of selu and does not dispatch any mul for the scale factor.
I am using pytorch 1.5
Hi,
Could you be more precise about what you want to do and how you implement it?
Are you registering new ops? Are you modifying pytorch source directly? How do you do that?
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:
static auto registerer = torch::RegisterOperators()
.op(torch::RegisterOperators::options().schema("aten::selu(Tensor self) -> Tensor")
.impl_unboxedOnlyKernel<at::Tensor(const at::Tensor &), &MyRegisteredOps::selu>(at::DispatchKey::NewDispatchKey)
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.
Hi,
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.
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 PrivateUse1 that is especially design to allow users to have their own backend.
thanks I will uplift to using PrivateUse1
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?
I’m not sure what you mean by that.
Can you give a concrete example?
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.
If the PrivateUse1 is listed after the autograd dispatch key would that result in selu being dispatched?
Ho so you’re trying to register new dispatch keys for operators that are already defined in core pytorch?
We use the term “leaf” in the autograd usually but it is unrelated here 
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.
cc @Sebastian_Messmer who will know better.
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. @Sebastian_Messmer
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.
How would I intercept the first call and do you mind explaining what would happen in the case I want to do Autograd?
I am also seeing a case where the forward function calls torch.flatten(x) but I don’t see the flatten
Thanks for your help!
The standard operator registration should work. Something like:
static auto registry = torch::RegisterOperators().op("aten::selu", torch::RegisterOperators::options()
.kernel(DispatchKey::PrivateUse1, &your_selu_implementation)
);
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.
Hi @almeetb, in pytorch dispatcher compound=True means we have a mathematical definition of selu that works for all backends and we just by default use that implementation. So even you register a kernel to selu it won’t get called, instead it’ll be called at eluwhich is a non-compound op.
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 https://github.com/pytorch/xla/blob/master/torch_xla/csrc/aten_autograd_ops.cpp#L18. You’ll need to override both forward and backward through torch::autograd::Function, and register that kernel to a XLAPreAutograd key.
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.
The recommended way for now is to not intercept compound ops and instead override leaf nodes like elu.