Merge pull request #6137 from kmaehashi/document-ufunc

Generate signature for ufunc documentation

cupy/__init__.py

@@ -925,3 +925,21 @@ def __getattr__(name):
         return getattr(_numpy, name)
 
     raise AttributeError(f"module 'cupy' has no attribute {name!r}")
+
+
+def _embed_signatures(dirs):
+    for name, value in dirs.items():
+        if isinstance(value, ufunc):
+            from cupy._core._kernel import _ufunc_doc_signature_formatter
+            value.__doc__ = (
+                _ufunc_doc_signature_formatter(value, name) +
+                '\n\n' + value._doc
+            )
+
+
+_embed_signatures(globals())
+_embed_signatures(fft.__dict__)
+_embed_signatures(linalg.__dict__)
+_embed_signatures(random.__dict__)
+_embed_signatures(sparse.__dict__)
+_embed_signatures(testing.__dict__)

cupy/_core/_kernel.pyx

@@ -1153,7 +1153,8 @@ cdef class ufunc:
         readonly tuple _params_with_where
         readonly dict _routine_cache
         readonly dict _kernel_memo
-        readonly object __doc__
+        readonly object _doc
+        public object __doc__
         readonly object __name__
         readonly object __module__
 
@@ -1170,11 +1171,13 @@ cdef class ufunc:
         self._out_ops = out_ops
         self._preamble = preamble
         self._loop_prep = loop_prep
+        self._doc = doc
         self.__doc__ = doc
         if default_casting is None:
             self._default_casting = 'same_kind'
         else:
             self._default_casting = default_casting
+
         if cutensor_op is not None and cuda_cutensor is not None:
             self._cutensor_op, self._cutensor_alpha, self._cutensor_gamma = (
                 getattr(cuda_cutensor, cutensor_op[0]),
@@ -1466,6 +1469,42 @@ cdef class ufunc:
             f'`{self.name}.reduceat` is not supported yet')
 
 
+def _ufunc_doc_signature_formatter(ufunc, name):
+    # Based on implementation in NumPy (numpy/_core/_internal.py)
+
+    # input arguments are simple
+    if ufunc.nin == 1:
+        in_args = 'x'
+    else:
+        in_args = ', '.join(f'x{i+1}' for i in range(ufunc.nin))
+
+    # output arguments are both keyword or positional
+    if ufunc.nout == 0:
+        out_args = ', /, out=()'
+    elif ufunc.nout == 1:
+        out_args = ', /, out=None'
+    else:
+        out_args = '[, {positional}], / [, out={default}]'.format(
+            positional=', '.join(
+                'out{}'.format(i+1) for i in range(ufunc.nout)),
+            default=repr((None,)*ufunc.nout)
+        )
+
+    # keyword only args depend on whether this is a gufunc
+    kwargs = (
+        f", casting='{ufunc._default_casting}'"
+        ", dtype=None"
+    )
+
+    # join all the parts together
+    return r'{name}({in_args}{out_args}, \*{kwargs})'.format(
+        name=name,
+        in_args=in_args,
+        out_args=out_args,
+        kwargs=kwargs
+    )
+
+
 cdef class _Op:
 
     def __init__(

docs/source/conf.py

@@ -17,6 +17,7 @@
 import os
 import sys
 
+import sphinx
 import cupy
 
 sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
@@ -498,6 +499,19 @@ def remove_array_api_module_docstring(app, what, name, obj, options, lines):
     if what == "module" and 'array_api' in name:
         del lines[:]
 
+
+def _patch_function_documenter():
+    # Monkeypatch FunctionDocumenter to let autosummary document ufuncs with
+    # `autofunction` instead of `autodata` so that signatures are documented.
+    documenter = sphinx.ext.autodoc.FunctionDocumenter
+    orig = documenter.can_document_member
+    def _can_document_member(member, *args, **kwargs):
+        if isinstance(member, cupy.ufunc):
+            return True
+        return orig(member, *args, **kwargs)
+    documenter.can_document_member = _can_document_member
+
+
 def fix_jit_callable_signature(
         app, what, name, obj, options, signature, return_annotation):
     if 'cupyx.jit' in name and callable(obj) and signature is None:
@@ -514,6 +528,8 @@ def fix_ndarray_signature(
     return (signature, return_annotation)
 
 def setup(app):
+    _patch_function_documenter()
+
     app.connect("autodoc-process-docstring", remove_array_api_module_docstring)
     app.connect("autodoc-process-signature", fix_jit_callable_signature)
     app.connect("autodoc-process-signature", fix_ndarray_signature)