pytorch · prasunanand · Aug 17, 2019 · Aug 19, 2019 · Aug 19, 2019 · Aug 19, 2019
diff --git a/docs/source/notes/extending.rst b/docs/source/notes/extending.rst
@@ -2,7 +2,7 @@ Extending PyTorch
 =================
 
 In this note we'll cover ways of extending :mod:`torch.nn`,
-:mod:`torch.autograd`, and writing custom C extensions utilizing our C
+:mod:`torch.autograd`, :mod:`torch`, and writing custom C extensions utilizing our C
 libraries.
 
 Extending :mod:`torch.autograd`
@@ -204,6 +204,285 @@ This is how a ``Linear`` module can be implemented::
                 self.in_features, self.out_features, self.bias is not None
             )
 
+Extending :mod:`torch`
+----------------------
+
+You can create custom types that emulate :class:`Tensor` by defining a custom
+class with methods that match :class:`Tensor`. But what if you want to be able
+to pass these types to functions like :func:`torch.add` in the top-level
+:mod:`torch` namespace that accept :class:`Tensor` operands?
+
+If your custom python type defines a method named ``__torch_function__``, PyTorch
+will invoke your ``__torch_function__`` implementation when an instance of your
+custom class is passed to a function in the :mod:`torch` namespace. This makes
+it possible to define custom implementations for any of the functions in the
+:mod:`torch` namespace which your ``__torch_function__`` implementation can call,
+allowing your users to make use of your custom type with existing PyTorch
+workflows that they have already written for :class:`Tensor`. This works with
+"duck" types that are unrelated to :class:`Tensor` as well as user-defined
+subclasses of :class:`Tensor`.
+
+Extending :mod:`torch` with a :class:`Tensor`-like type
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. note:: This functionality is inspired by the NumPy ``__array_function__``
+          protocol. See `the NumPy documentation
+          <https://docs.scipy.org/doc/numpy/user/basics.dispatch.html#basics-dispatch>`_
+          and `NEP-0018
+          <https://numpy.org/neps/nep-0018-array-function-protocol.html>`_ for
+          more details.
+
+To make this concrete, let's begin with a simple example that illustrates the
+API dispatch mechanism. We'll create a custom type that represents a 2D scalar
+tensor, parametrized by the order ``N`` and value along the diagonal entries,
+``value``::
+
+     class ScalarTensor(object):
+        def __init__(self, N, value):
+            self._N = N
+            self._value = value
+
+        def __repr__(self):
+            return "DiagonalTensor(N={}, value={})".format(self._N, self._value)
+
+        def tensor(self):
+            return self._value * torch.eye(self._N)
+
+This first iteration of the design isn't very useful. The main functionality of
+``ScalarTensor`` is to provide a more compact string representation of a scalar
+tensor than in the base tensor class::
+
+  >>> d = ScalarTensor(5, 2)
+  >>> d
+  ScalarTensor(N=5, value=2)
+  >>> d.tensor()
+  tensor([[2., 0., 0., 0., 0.],
+          [0., 2., 0., 0., 0.],
+          [0., 0., 2., 0., 0.],
+          [0., 0., 0., 2., 0.],
+          [0., 0., 0., 0., 2.]])
+
+If we try to use this object with the :mod:`torch` API, we will run
+into issues::
+
+  >>> import torch
+  >>> torch.mean(d)
+  TypeError: mean(): argument 'input' (position 1) must be Tensor, not ScalarTensor
+
+Adding a ``__torch_function__`` implementation to ``ScalarTensor`` makes it
+possible for the above operation to succeed. Let's re-do our implementation,
+this time adding a ``__torch_function__`` implementation::
+
+  HANDLED_FUNCTIONS = {}
+  class ScalarTensor(object):
+      def __init__(self, N, value):
+          self._N = N
+          self._value = value
+
+      def __repr__(self):
+          return "DiagonalTensor(N={}, value={})".format(self._N, self._value)
+
+      def tensor(self):
+          return self._value * torch.eye(self._N)
+
+      def __torch_function__(self, func, args=(), kwargs=None):
+          if kwargs is None:
+              kwargs = {}
+          if func not in HANDLED_FUNCTIONS:
+              return NotImplemented
+          return HANDLED_FUNCTIONS[func](*args, **kwargs)
+
+The ``__torch_function__`` method takes three arguments: ``func``, a reference to
+the torch API function that is being overrided, ``args``, the tuple of arguments
+passed to the function, and ``kwargs``, the dict of keyword arguments passed to
+the function. It uses a global dispatch stable named ``HANDLED_FUNCTIONS`` to
+store custom implementations. The keys of this dictionary are functions in the
+``torch`` namespace and the values are implementations for ``ScalarTensor``.
+
+.. note:: Using a global dispatch table is not a mandated part of the
+          ``__torch_function__`` API, it is just a useful design pattern for
+          structuring your override implementations.
+
+This class definition isn't quite enough to make ``torch.mean`` do the right
+thing when we pass it a ``ScalarTensor`` -- we also need to define an
+implementation for ``torch.mean`` for ``ScalarTensor`` operands and add the
+implementation to the ``HANDLED_FUNCTIONS`` dispatch table dictionary. One way
+of doing this is to define a decorator::
+
+  import functools
+  def implements(torch_function):
+      """Register a torch function override for ScalarTensor"""
+      @functools.wraps(torch_function)
+      def decorator(func):
+          HANDLED_FUNCTIONS[torch_function] = func
+          return func
+      return decorator
+
+which can be applied to the implementation of our override::
+
+  @implements(torch.mean)
+  def mean(input):
+      return float(input._value) / input._N
+
+With this change we can now use ``torch.mean`` with ``ScalarTensor``::
+
+  >>> d = ScalarTensor(5, 2)
+  >>> torch.mean(d)
+  0.4
+
+Of course ``torch.mean`` is an example of the simplest kind of function to
+override since it only takes one operand. We can use the same machinery to
+override a function that takes more than one operand, any one of which might be
+a tensor or tensor-like that defines ``__torch_function__``, for example for
+:func:`torch.add`::
+
+  def ensure_tensor(data):
+      if isinstance(data, ScalarTensor):
+          return data.tensor()
+      return torch.as_tensor(data)
+
+  @implements(torch.add)
+  def add(input, other):
+     try:
+         if input._N == other._N:
+             return ScalarTensor(input._N, input._value + other._value)
+         else:
+             raise ValueError("Shape mismatch!")
+     except AttributeError:
+         return torch.add(ensure_tensor(input), ensure_tensor(other))
+
+This version has a fast path for when both operands are ``ScalarTensor``
+instances and also a slower path which degrades to converting the data to
+tensors when either operand is not a ``ScalarTensor``. That makes the override
+function correctly when either operand is a ``ScalarTensor`` or a regular
+:class:`Tensor`::
+
+  >>> s = ScalarTensor(2, 2)
+  >>> torch.add(s, s)
+  DiagonalTensor(N=2, value=4)
+  >>> t = torch.tensor([[1, 1,], [1, 1]])
+  >>> torch.add(s, t)
+  tensor([[3., 1.],
+          [1., 3.]])
+
+Note that our implementation of ``add`` does not take ``alpha`` or ``out`` as
+keyword arguments like :func:`torch.add` does::
+
+  >>> torch.add(s, s, alpha=2)
+  TypeError: add() got an unexpected keyword argument 'alpha'
+
+For speed and flexibility the ``__torch_function__`` dispatch mechanism does not
+check that the signature of an override function matches the signature of the
+function being overrided in the :mod:`torch` API. For some applications ignoring
+optional arguments would be fine but to ensure full compatibility with
+:class:`Tensor`, user implementations of torch API functions should take care to
+exactly emulate the API of the function that is being overrided.
+
+Functions in the :mod:`torch` API that do not have explicit overrides will
+return ``NotImplemented`` from ``__torch_function__``. If all operands with
+``__torch_function__`` defined on them return ``NotImplemented``, PyTorch will
+raise a ``TypeError``. This means that most of the time operations that do not
+have explicit overrides for a type will raise a ``TypeError`` when an instance
+of such a type is passed::
+
+  >>> torch.mul(s, 3)
+  TypeError: no implementation found for 'torch.mul' on types that
+  implement __torch_function__: [ScalarTensor]
+
+In practice this means that if you would like to implement your overrides using
+a ``__torch_function__`` implementation along these lines, you will need to
+explicitly implement the full :mod:`torch` API or the entire subset of the API
+that you care about for your use case. This may be a tall order as the full
+:mod:`torch` API is quite extensive.
+
+Another option is to not return ``NotImplemented`` for operations that are not
+handled but to instead pass a :class:`Tensor` to the original :mod:`torch`
+function when no override is available. For example, if we change our
+implementation of ``__torch_function__`` for ``ScalarTensor`` to the one below::
+
+  def __torch_function__(self, func, args=(), kwargs=None):
+      if kwargs is None:
+          kwargs = {}
+      if func not in HANDLED_FUNCTIONS:
+          args = [a.tensor() if hasattr(a, 'tensor') else a for a in args]
+          return func(*args, **kwargs)
+      return HANDLED_FUNCTIONS[func](*args, **kwargs)
+
+Then :func:`torch.mul` will work correctly, although the return type will always
+be a :class:`Tensor` rather than a :class:`ScalarTensor`, even if both operands
+are :class:`ScalarTensor` instances::
+
+  >>> s = ScalarTensor(2, 2)
+  >>> torch.mul(s, s)
+  tensor([[4., 0.],
+          [0., 4.]])
+
+Also see the ``MetadataTensor`` example below for another variation on this
+pattern but instead always returns a ``MetadataTensor`` to propagate metadata
+through operations in the :mod:`torch` API.
+
+Extending :mod:`torch` with a :class:`Tensor` wrapper type
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Another useful case is a type that wraps a :class:`Tensor`, either as an
+attribute or via subclassing. Below we implement a special case of this sort of
+type, a ``MetadataTensor`` that attaches a dictionary of metadata to a
+:class:`Tensor` that is propagated through :mod:`torch` operations. Since this
+is a generic sort of wrapping for the full :mod:`torch` API, we do not need to
+individually implement each override so we can make the ``__torch_function__``
+implementation more permissive about what operations are allowed::
+
+  class MetadataTensor(object):
+      def __init__(self, data, metadata=None, **kwargs):
+          self._t = torch.as_tensor(data, **kwargs)
+          self._metadata = metadata
+
+      def __repr__(self):
+          return "Metadata:\n{}\n\ndata:\n{}".format(self._metadata, self._t)
+
+      def __torch_function__(self, func, args=(), kwargs=None):
+          if kwargs is None:
+              kwargs = {}
+          args = [a._t if hasattr(a, '_t') else a for a in args]
+          ret = func(*args, **kwargs)
+          return MetadataTensor(ret, metadata=self._metadata)
+
+This simple implementation won't necessarily work with every function in the
+:mod:`torch` API but it is good enough to capture most common operations::
+
+  >>> metadata = {'owner': 'Ministry of Silly Walks'}
+  >>> m = MetadataTensor([[1, 2], [3, 4]], metadata=metadata)
+  >>> t = torch.tensor([[1, 2], [1, 2]]])
+  >>> torch.add(t, m)
+  Metadata:
+  {'owner': 'Ministry of Silly Walks'}
+
+  data:
+  tensor([[2, 4],
+          [4, 6]])
+  >>> torch.mul(t, m)
+  Metadata:
+  {'owner': 'Ministry of Silly Walks'}
+
+  data:
+  tensor([[1, 4],
+          [3, 8]])
+
+Operations on multiple types that define ``__torch_function__``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+It is possible to use the torch API with multiple distinct types that each have
+a ``__torch_function__`` implementation, but special care must be taken. In such
+a case the rules are:
+
+* The dispatch operation gathers all distinct implementations of
+  ``__torch_function__`` for each operand and calls them in order: subclasses
+  before superclasses, and otherwise left to right in the operator expression.
+* If any value other than ``NotImplemented`` is returned, that value is
+  returned as the result. Implementations can register that they do not
+  implement an operation by returning ``NotImplemented``.
+* If all of the ``__torch_function__`` implementations return
+  ``NotImplemented``, PyTorch raises a ``TypeError``.
 
 Writing custom C++ extensions
 -----------------------------

diff --git a/test/onnx/expect/TestOperators.test_frobenius_norm.expect b/test/onnx/expect/TestOperators.test_frobenius_norm.expect
@@ -3,8 +3,8 @@ producer_name: "pytorch"
 producer_version: "1.3"
 graph {
   node {
-    input: "x"
-    input: "x"
+    input: "0"
+    input: "0"
     output: "1"
     name: "Mul_0"
     op_type: "Mul"
@@ -34,7 +34,7 @@ graph {
   }
   name: "torch-jit-export"
   input {
-    name: "x"
+    name: "0"
     type {
       tensor_type {
         elem_type: 1

diff --git a/test/onnx/expect/TestOperators.test_meshgrid.expect b/test/onnx/expect/TestOperators.test_meshgrid.expect
@@ -17,7 +17,7 @@ graph {
     }
   }
   node {
-    input: "x"
+    input: "0"
     input: "3"
     output: "4"
     name: "Reshape_1"
@@ -38,7 +38,7 @@ graph {
     }
   }
   node {
-    input: "y"
+    input: "1"
     input: "5"
     output: "6"
     name: "Reshape_3"
@@ -59,7 +59,7 @@ graph {
     }
   }
   node {
-    input: "z"
+    input: "2"
     input: "7"
     output: "8"
     name: "Reshape_5"
@@ -221,7 +221,7 @@ graph {
   }
   name: "torch-jit-export"
   input {
-    name: "x"
+    name: "0"
     type {
       tensor_type {
         elem_type: 1
@@ -234,7 +234,7 @@ graph {
     }
   }
   input {
-    name: "y"
+    name: "1"
     type {
       tensor_type {
         elem_type: 1
@@ -247,7 +247,7 @@ graph {
     }
   }
   input {
-    name: "z"
+    name: "2"
     type {
       tensor_type {
         elem_type: 1

diff --git a/test/onnx/expect/TestOperators.test_unique.expect b/test/onnx/expect/TestOperators.test_unique.expect
@@ -3,7 +3,7 @@ producer_name: "pytorch"
 producer_version: "1.3"
 graph {
   node {
-    input: "x"
+    input: "0"
     output: "1"
     output: "2"
     output: "3"
@@ -23,7 +23,7 @@ graph {
   }
   name: "torch-jit-export"
   input {
-    name: "x"
+    name: "0"
     type {
       tensor_type {
         elem_type: 1

diff --git a/test/run_test.py b/test/run_test.py
@@ -62,6 +62,7 @@
     'type_promotion',
     'jit_disabled',
     'function_schema',
+    'overrides',
 ]
 
 # skip < 3.3 because mock is added in 3.3 and is used in rpc_spawn