Skip to content

Add section about C++ classes as op args #880

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Mar 10, 2020
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
44 changes: 44 additions & 0 deletions advanced_source/torch_script_custom_classes.rst
Original file line number Diff line number Diff line change
Expand Up @@ -520,6 +520,50 @@ now run successfully:
$ python ../export_attr.py
testing

Defining Custom Operators that Take C++ Classes as Arguments
------------------------------------------------------------

Once you've defined a custom C++ class, you can also use that class
as an argument to custom operator (i.e. free functions). Here's an
example of how to do that:

.. code-block:: cpp

std::string take_an_instance(const c10::intrusive_ptr<Stack<std::string>>& instance) {
return instance->pop();
}

static auto instance_registry = torch::RegisterOperators().op(
torch::RegisterOperators::options()
.schema(
"foo::take_an_instance(__torch__.torch.classes.Stack x) -> str Y")
.catchAllKernel<decltype(take_an_instance), &take_an_instance>());

Refer to the `custom op tutorial <https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html>`_
for more details on the registration API.

Once this is done, you can use the op like the following example:

.. code-block:: python

class TryCustomOp(torch.nn.Module):
def __init__(self):
super(TryCustomOp, self).__init__()
self.f = torch.classes.Stack(["foo", "bar"])

def forward(self) -> str:
return torch.ops._TorchScriptTesting.take_an_instance(self.f)

.. note::

Registration of an operator that takes a C++ class as an argument requires that
the custom class has already been registered. This is fine if your op is
registered after your class in a single compilation unit, however, if your
class is registered in a separate compilation unit from the op you will need
to enforce that dependency. One way to do this is to wrap the class registration
in a `Meyer's singleton <https://stackoverflow.com/q/1661529>`_, which can be
called from the compilation unit that does the operator registration.

Conclusion
----------

Expand Down