Update

Author: zou3519

advanced_source/cpp_custom_ops.rst

@@ -167,13 +167,38 @@ for more details).
 
 .. code-block:: python
 
+  # Important: the C++ custom operator definitions should be loaded first
+  # before calling ``torch.library`` APIs that add registrations for the
+  # C++ custom operator(s). The following import loads our
+  # C++ custom operator definitions.
+  # See the next section for more details.
+  from . import _C
+
   @torch.library.register_fake("extension_cpp::mymuladd")
   def _(a, b, c):
       torch._check(a.shape == b.shape)
       torch._check(a.dtype == torch.float)
       torch._check(b.dtype == torch.float)
       torch._check(a.device == b.device)
       return torch.empty_like(a)
+
+How to set up hybrid Python/C++ registration
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+In this tutorial, we defined a custom operator in C++, added CPU/CUDA
+implementations in C++, and added ``FakeTensor`` kernels and backward formulas
+in Python. The order in which these registrations are loaded (or imported)
+matters (importing in the wrong order will lead to an error).
+
+To use the custom operator with hybrid Python/C++ registrations, we must
+first load the C++ library that holds the custom operator definition
+and then call the ``torch.library`` registration APIs. This can happen in one
+of two ways:
+
+1. If you're following this tutorial, importing the Python C extension module
+   we created will load the C++ custom operator definitions.
+2. If your C++ custom operator is located in a shared library object, you can
+   also use ``torch.ops.load_library("/path/to/library.so")`` to load it.
+
 
 How to add training (autograd) support for an operator
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

advanced_source/python_custom_ops.py

@@ -74,7 +74,7 @@ def f(img):
 # do two things:
 #
 # 1. wrap the function into a PyTorch custom operator.
-# 2. add a "FakeTensor kernel" (aka "meta kernel") to the operator.
+# 2. add a "``FakeTensor`` kernel" (aka "meta kernel") to the operator.
 #    Given the metadata (e.g. shapes)
 #    of the input Tensors, this function says how to compute the metadata
 #    of the output Tensor(s).
@@ -84,22 +84,22 @@ def f(img):
 
 # Use torch.library.custom_op to define a new custom operator.
 # If your operator mutates any input Tensors, their names must be specified
-# in the mutates_args argument.
+# in the ``mutates_args`` argument.
 @torch.library.custom_op("mylib::crop", mutates_args=())
 def crop(pic: torch.Tensor, box: Sequence[int]) -> torch.Tensor:
     img = to_pil_image(pic.cpu())
     cropped_img = img.crop(box)
     return (pil_to_tensor(cropped_img) / 255.).to(pic.device, pic.dtype)
 
-# Use register_fake to add a FakeTensor kernel for the operator
+# Use register_fake to add a ``FakeTensor`` kernel for the operator
 @crop.register_fake
 def _(pic, box):
     channels = pic.shape[0]
     x0, y0, x1, y1 = box
     return pic.new_empty(channels, y1 - y0, x1 - x0)
 
 ######################################################################
-# After this, ``crop`` now works whout graph breaks:
+# After this, ``crop`` now works without graph breaks:
 
 @torch.compile(fullgraph=True)
 def f(img):
@@ -214,7 +214,7 @@ def numpy_sin(input: torch.Tensor, output: torch.Tensor) -> None:
 
 ######################################################################
 # Because the operator doesn't return anything, there is no need to register
-# a FakeTensor kernel (meta kernel) to get it to work with ``torch.compile``.
+# a ``FakeTensor`` kernel (meta kernel) to get it to work with ``torch.compile``.
 
 @torch.compile(fullgraph=True)
 def f(x):

en-wordlist.txt

@@ -177,6 +177,7 @@ OpenSlide
 Opset
 Ornstein
 PIL
+PIL's
 PPO
 PatchPredictor
 PennFudan
@@ -443,6 +444,7 @@ multithreading
 namespace
 natively
 ndarrays
+nightlies
 num
 numericalize
 numpy