Changes to autograd/custom functions to handle optional arguments (#54270)
Summary:
Small changes to autograd to support optional Tensor values.
On MLC device, we use Autograd Custom Functions to override the autograd engine for a specific operation. We do something like:
```
at::Tensor AtenMLCAutogradTypeDefault::abs(const at::Tensor & self) {
torch_mlc::mlclogger() << "MLC bridge autograd MLC : abs" << std::endl;
torch_mlc::AutoNonAtenMLCAutogradTypeDefault guard(true);
return MLCAbsFunction::apply(self);
}
TORCH_LIBRARY_IMPL(aten, AutogradMLC, m) {
m.impl("abs", static_cast<at::Tensor (*)(const at::Tensor &)>(&AtenMLCAutogradTypeDefault::abs));
}
```
What I noticed is that the existing code does not always work for optional Tensor types. This PR fixes it. Let me know if you have a better way to deal with this issue.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/54270
Reviewed By: ejguan
Differential Revision: D27171623
Pulled By: albanD
fbshipit-source-id: 3aa8d59ee8da3cc943ad5e73521c2755d1ff2341
