__torch_function__ overrides for torch.functional and torch.nn.functional (pytorch#32799)

Summary:
This adds `__torch_function__` support for all functions in `torch.functional` and `torch.nn.functional`.

The changes to C++ code and codegen scripts are to facilitate adding `__torch_function__` support for the native functions in `torch._C._nn`. Note that I moved the `handle_torch_function` C++ function to a header that both `python_torch_functions.cpp` and `python_nn_functions.cpp` include. The changes to `python_nn_functions.cpp` mirror the changes I made to `python_torch_functions.cpp` when `__torch_function__` support was first added in pytorch#27064. Due to the somewhat different way the `torch._C` and `torch._C._nn` namespaces are initialized I needed to create a new static reference to the `torch._C._nn` namespace (`THPNNVariableFunctions`). I'm not sure if that is the best way to do this. In principle I could import these namespaces in each kernel and avoid the global variable but that would have a runtime cost.

I added `__torch_function__` support to the Python functions in `torch.nn.functional` following the approach in pytorch#32194.

I re-enabled the test that checks if all functions in the `torch` namespace are explicitly tested for `__torch_function__` support. I also generalized the check to work for `torch.functional` and `torch.nn.functional` as well. This test was explicitly disabled in pytorch#30730 and I'm happy to disable it again if you think that's appropriate. I figured now was as good a time as any to try to re-enable it.

Finally I adjusted the existing torch API tests to suppress deprecation warnings and add keyword arguments used by some of the code in `torch.nn.functional` that were missed when I originally added the tests in pytorch#27064.
Pull Request resolved: pytorch#32799

Differential Revision: D19956809

Pulled By: ezyang

fbshipit-source-id: 40d34e0109cc4b9f3ef62f409d2d35a1d84e3d22

Author: ngoldbaum

test/test_overrides.py

@@ -82,6 +82,11 @@
     'div(Tensor, Scalar)', 'div_(Tensor, Scalar)',
 ]
 
+NATIVE_NAMESPACE_MAPPING = {
+    "torch": "THPVariableFunctionsModule",
+    "torch.nn": "THPNNVariableFunctionsModule"
+}
+
 def should_generate_python_binding(declaration):
     name = declaration['name']
     for pattern in SKIP_PYTHON_BINDINGS:
@@ -121,7 +126,8 @@ def gen_py_variable_methods(out, declarations, template_path):
 
     py_variable_methods = get_py_variable_methods(declarations)
 
-    env = create_python_bindings(py_variable_methods, is_python_method=True, is_module=False)
+    env = create_python_bindings(py_variable_methods, is_python_method=True, module=None)
+
     write(out, 'python_variable_methods.cpp', PY_VARIABLE_METHODS_CPP, env)
 
 
@@ -145,7 +151,8 @@ def gen_py_nn_functions(out, declarations, template_path):
 
     py_nn_functions = get_py_nn_functions(declarations)
 
-    env = create_python_bindings(py_nn_functions, is_python_method=False, is_module=True)
+    env = create_python_bindings(py_nn_functions, is_python_method=False, module="torch.nn")
+
     write(out, 'python_nn_functions.cpp', PY_NN_FUNCTIONS_CPP, env)
 
 
@@ -170,7 +177,8 @@ def gen_py_torch_functions(out, declarations, template_path):
 
     py_torch_functions = get_py_torch_functions(declarations)
 
-    env = create_python_bindings(py_torch_functions, is_python_method=False, is_module=False)
+    env = create_python_bindings(py_torch_functions, is_python_method=False, module="torch")
+
     write(out, 'python_torch_functions.cpp', PY_TORCH_FUNCTIONS_CPP, env)
 
 
@@ -181,17 +189,17 @@ def group_declarations_by_op_name(declarations):
     return groups
 
 
-def create_python_bindings(python_functions, is_python_method, is_module):
+def create_python_bindings(python_functions, is_python_method, module):
     """Generates Python bindings to ATen functions"""
     py_methods = []
     py_method_defs = []
     py_forwards = []
 
     for name in sorted(python_functions.keys()):
         overload_decls = python_functions[name]
-        py_methods.append(method_impl(name, overload_decls, is_python_method, is_module))
-        py_method_defs.append(method_def(name, overload_decls, is_python_method, is_module))
-        py_forwards.extend(forward_decls(name, overload_decls, is_python_method, is_module))
+        py_methods.append(method_impl(name, overload_decls, is_python_method, module))
+        py_method_defs.append(method_def(name, overload_decls, is_python_method, module))
+        py_forwards.extend(forward_decls(name, overload_decls, is_python_method, module))
 
     return {
         'py_forwards': py_forwards,
@@ -713,7 +721,6 @@ def get_field_name(x):
                 return x['field_name']
         return [get_field_name(x) for x in returns]
 
-
 PY_NAMEDTUPLE_FIELDSDEF = CodeTemplate("""\
 static PyStructSequence_Field ${fieldsname}[] = { ${fields,} {nullptr} };
 """)
@@ -840,19 +847,19 @@ def is_noarg_binding(overloads):
 
 """)
 
-TORCH_FUNCTION_CHECK = """\
-if (_r.has_torch_function()) {
-  return handle_torch_function(_r, args, kwargs, THPVariableFunctions);
+TORCH_FUNCTION_CHECK = CodeTemplate("""\
+if(_r.has_torch_function()) {
+  return handle_torch_function(_r, args, kwargs, ${namespace}, ${modulename});
 }
-"""
+""")
 
 # NOTE: we type the unpacked self as Tensor not Variable to avoid return type
 # discrepancies on method resolution (e.g. Variable::detach_ returns void
 # rather than Tensor &)
 UNPACK_SELF = "Tensor& self = reinterpret_cast<THPVariable*>(self_)->cdata;"
 
 
-def method_impl(name, declarations, is_python_method, is_module):
+def method_impl(name, declarations, is_python_method, module):
     """
     Generate a python binding for all overloads of an op.
     """
@@ -897,8 +904,11 @@ def method_impl(name, declarations, is_python_method, is_module):
     else:
         template = PY_VARIABLE_METHOD_VARARGS
 
-    if not is_module and not is_python_method:
-        check_has_torch_function = TORCH_FUNCTION_CHECK
+    if module:
+        check_has_torch_function = TORCH_FUNCTION_CHECK.substitute(
+            namespace=NATIVE_NAMESPACE_MAPPING[module],
+            modulename='"' + module + '"',
+        )
     else:
         check_has_torch_function = ''
 
@@ -931,8 +941,8 @@ def method_impl(name, declarations, is_python_method, is_module):
 """)
 
 
-def forward_decls(name, declarations, is_python_method, is_module):
-    if is_module or is_python_method:
+def forward_decls(name, declarations, is_python_method, module):
+    if is_python_method:
         return []
 
     if is_noarg_binding(declarations):
@@ -979,7 +989,7 @@ def forward_decls(name, declarations, is_python_method, is_module):
 {"${name}", (PyCFunction)${pycfunc_voidcast}${pycname}, ${flags}, NULL},""")
 
 
-def method_def(name, declarations, is_python_method, is_module):
+def method_def(name, declarations, is_python_method, module):
     """
     Generate method def entry.
     """
@@ -992,7 +1002,7 @@ def method_def(name, declarations, is_python_method, is_module):
         pycfunc_voidcast = '(void(*)(void))'
         flags = 'METH_VARARGS | METH_KEYWORDS'
 
-    if not is_module and not is_python_method:
+    if module == "torch":
         flags += ' | METH_STATIC'
 
     if name in BINARY_OP_NAMES:

tools/autograd/gen_python_functions.py

@@ -53,14 +53,18 @@ static PyObject * THPVariable__parse_to(PyObject* module, PyObject* args, PyObje
   END_HANDLE_TH_ERRORS
 }
 
-${py_methods}
+// generated forward declarations start here
+
+${py_forwards}
 
 static PyMethodDef nn_functions[] = {
   {"_parse_to", (PyCFunction)(void(*)(void))THPVariable__parse_to, METH_VARARGS | METH_KEYWORDS, nullptr},
   ${py_method_defs}
   {NULL}
 };
 
+static PyObject* THPNNVariableFunctionsModule = NULL;
+
 void initNNFunctions(PyObject* module) {
 #if PY_MAJOR_VERSION == 2
   PyObject* nn = Py_InitModule("torch._C._nn", nn_functions);
@@ -75,6 +79,7 @@ void initNNFunctions(PyObject* module) {
   };
   PyObject* nn = PyModule_Create(&def);
 #endif
+  THPNNVariableFunctionsModule = nn;
   if (!nn) {
     throw python_error();
   }
@@ -84,4 +89,8 @@ void initNNFunctions(PyObject* module) {
   }
 }
 
+// generated methods start here
+
+${py_methods}
+
 }} // namespace torch::autograd

tools/autograd/templates/python_nn_functions.cpp

@@ -450,88 +450,19 @@ static PyTypeObject THPVariableFunctions = {
   0                                      /* tp_new */
 };
 
+static PyObject* THPVariableFunctionsModule = NULL;
+
 void initTorchFunctions(PyObject* module) {
   if (PyType_Ready(&THPVariableFunctions) < 0) {
     throw python_error();
   }
   Py_INCREF(&THPVariableFunctions);
-  if (PyModule_AddObject(module, "_VariableFunctions", (PyObject*)&THPVariableFunctions) < 0) {
-    throw python_error();
-  }
-}
-
-/*
- *
- * Calls __torch_function__ on the overloaded arguments to a torch API
- * function in order of precedence, returning the first result that is
- * not NotImplemented. If all arguments return NotImplemented, raises a
- * TypeError.
- *
- * Assumes overloaded_args has at least one entry. All entries must have
- * a __torch_function__ attribute that resolves to a callable that
- * accepts a torch API function, arguments, and keyword arguments for
- * the torch API function.
- *
- * It is sufficient to call PythonArgs::has_torch_function before
- * calling this function to verify that there are valid arguments
- * present. If that is not done then special care must be taken to
- * ensure there are arguments that are overloaded with
- * __torch_function__.
- *
- * See torch._overrides._implement_torch_function for the equivalent
- * code in the pure-python implementation.
- *
- * 'r' is a parsed PythonArgs instance, returned from
- * PythonArgParser::parse.
- *
- * 'args' is a reference to the python tuple of arguments to the torch
- * API function.
- *
- * 'kwargs' is a reference to the python dict of keyword arguments to
- * the torch API function.
- *
- * 'torch_api' is a reference to python torch API namespace.
- *
- */
-
-PyObject* handle_torch_function(PythonArgs &r, PyObject* args, PyObject* kwargs, PyTypeObject &torch_api) {
-  py::object torch_api_function = PyObject_FastGetAttrString((PyObject*)&torch_api, const_cast<char*>(r.get_func_name().data()));
-  TORCH_INTERNAL_ASSERT(torch_api_function.ptr() != NULL, "torch API function must exist");
-  py::object ret;
-  for (auto &arg : r.signature.overloaded_args) {
-    py::object torch_function = PyObject_FastGetAttrString(arg.ptr(), "__torch_function__");
-    ret = py::reinterpret_steal<py::object>(PyObject_CallFunctionObjArgs(torch_function.ptr(), torch_api_function.ptr(), args, kwargs, NULL));
-    if (ret.ptr() != Py_NotImplemented) {
-      // Return the reference to the result. This also covers the case where ret
-      // is NULL and __torch_function__ raised an exception, which we throw below
-      break;
-    }
-  }
-  if (ret.ptr() == nullptr) {
-    // if an exception occurred in a user's implementation of
-    // __array_function__, throw it
-    throw python_error();
-  }
-  else if (ret.ptr() == Py_NotImplemented) {
-    // all __torch_function__ implementations in overloaded_args
-    // returned NotImplemented, so we raise a TypeError.
-    std::stringstream ss;
-    ss << "no implementation found for 'torch." << r.get_func_name()
-       << "' on types that implement __torch_function__: [";
-    for (auto &arg : r.signature.overloaded_args) {
-      ss << arg.ptr()->ob_type->tp_name;
-      if (!arg.is(r.signature.overloaded_args.back())) {
-        ss << ", ";
-      }
-      else {
-        ss << "]";
-      }
-    }
-    const std::string& tmp = ss.str();
-    PyErr_SetString(PyExc_TypeError, tmp.c_str());
+  // PyType_GenericNew returns a new reference
+  THPVariableFunctionsModule = PyType_GenericNew(&THPVariableFunctions, Py_None, Py_None);
+  // PyModule_AddObject steals a reference
+  if (PyModule_AddObject(module, "_VariableFunctions", THPVariableFunctionsModule) < 0) {
     throw python_error();
   }
-  return ret.release().ptr();
 }
 
 // generated methods start here
@@ -549,7 +480,7 @@ static PyObject * THPVariable_nonzero(PyObject* self, PyObject* args, PyObject*
   auto r = parser.parse(args, kwargs, parsed_args);
 
   if(r.has_torch_function()){
-    return handle_torch_function(r, args, kwargs, THPVariableFunctions);
+    return handle_torch_function(r, args, kwargs, THPVariableFunctionsModule, "torch");
   }
 
   if (r.idx == 0) {
@@ -579,7 +510,7 @@ static PyObject * THPVariable_numel(PyObject* self_, PyObject* args, PyObject* k
   auto r = parser.parse(args, kwargs, parsed_args);
 
   if(r.has_torch_function()){
-    return handle_torch_function(r, args, kwargs, THPVariableFunctions);
+    return handle_torch_function(r, args, kwargs, THPVariableFunctionsModule, "torch");
   }
 
   if (r.idx == 0) {

tools/autograd/templates/python_torch_functions.cpp

@@ -131,6 +131,46 @@ FunctionParameter::FunctionParameter(const std::string& fmt, bool keyword_only)
   }
 }
 
+auto handle_torch_function(PythonArgs &r, PyObject* args, PyObject* kwargs, PyObject* torch_api, const char* module_name) -> PyObject* {
+  py::object torch_api_function = PyObject_FastGetAttrString(torch_api, (char*)r.get_func_name().c_str());
+  TORCH_INTERNAL_ASSERT(torch_api_function.ptr() != nullptr, "torch API function must exist");
+  py::object ret;
+  for (auto &arg : r.signature.overloaded_args) {
+    py::object torch_function = PyObject_FastGetAttrString(arg.ptr(), "__torch_function__");
+    ret = py::reinterpret_steal<py::object>(PyObject_CallFunctionObjArgs(torch_function.ptr(), torch_api_function.ptr(), args, kwargs, NULL));
+    if (ret.ptr() != Py_NotImplemented) {
+      // Return the reference to the result. This also covers the case where ret
+      // is NULL and __torch_function__ raised an exception, which we throw below
+      break;
+    }
+  }
+  if (ret.ptr() == nullptr) {
+    // if an exception occurred in a user's implementation of
+    // __array_function__, throw it
+    throw python_error();
+  }
+  else if (ret.ptr() == Py_NotImplemented) {
+    // all __torch_function__ implementations in overloaded_args
+    // returned NotImplemented, so we raise a TypeError.
+    std::stringstream ss;
+    ss << "no implementation found for '" << module_name << "." << r.get_func_name()
+       << "' on types that implement __torch_function__: [";
+    for (auto &arg : r.signature.overloaded_args) {
+      ss << arg.ptr()->ob_type->tp_name;
+      if (!arg.is(r.signature.overloaded_args.back())) {
+        ss << ", ";
+      }
+      else {
+        ss << "]";
+      }
+    }
+    const std::string& tmp = ss.str();
+    PyErr_SetString(PyExc_TypeError, tmp.c_str());
+    throw python_error();
+  }
+  return ret.release().ptr();
+}
+
 /*
  *  obj has a __torch_function__ implementation and may either be a
  *  subclass of Tensor or a Tensor-like duck type. We may need to

torch/csrc/utils/python_arg_parser.cpp

@@ -676,4 +676,47 @@ static auto check_has_torch_function(PyObject* obj) -> bool
   return false;
 }
 
+/*
+ *
+ * Handle __torch_function__ overrides if we know that there are overloaded
+ * arguments.  All objects stored in r.overloaded_args must have a
+ * __torch_function__ implementation and the arguments must be ordered in order
+ * of precedence. Precedence goes from left to right in the order of the
+ * signature of the function the overloaded arguments were passed to, except
+ * subclasses are always considered before superclasses.
+ *
+ * If the result of calling __torch_function__ is NotImplemented, the
+ * next implementation in the precedence order is called. If all
+ * arguments return NotImplemented from their __torch_function__
+ * implementation, a TypeError is raised in Python.
+ *
+ * Assumes overloaded_args has at least one entry. All entries must have
+ * a __torch_function__ attribute that resolves to a callable that
+ * accepts a torch API function, a tuple of arguments, and a dict of
+ * keyword arguments for the torch API function.
+ *
+ * It is sufficient to call PythonArgs::has_torch_function before
+ * calling this function to verify that there are valid arguments
+ * present. If that is not done then special care must be taken to
+ * ensure there are arguments that are overloaded with
+ * __torch_function__.
+ *
+ * See torch._overrides.handle_torch_function for the equivalent
+ * code in the pure-python implementation.
+ *
+ * 'r' is a parsed PythonArgs instance, returned from
+ * PythonArgParser::parse.
+ *
+ * 'args' is a reference to the python tuple of arguments to the torch
+ * API function.
+ *
+ * 'kwargs' is a reference to the python dict of keyword arguments to
+ * the torch API function.
+ *
+ * 'torch_api' is a reference to a python torch API namespace.
+ *
+ */
+
+auto handle_torch_function(PythonArgs &r, PyObject* args, PyObject* kwargs, PyObject* torch_api, const char* module_name) -> PyObject*;
+
 } // namespace torch