doc(hansbug): update doc for numpy part

docs/source/conf.py

@@ -103,6 +103,8 @@ release = __VERSION__
 extensions = [
     'sphinx.ext.autodoc',
     'sphinx.ext.doctest',
+    'sphinx.ext.coverage',
+    'sphinx.ext.autosummary',
     'sphinx.ext.mathjax',
     'sphinx.ext.ifconfig',
     'sphinx.ext.viewcode',

treetensor/numpy/funcs.py

 import builtins
+from typing import List
 
 import numpy as np
 from treevalue import func_treelize as original_func_treelize
 
 from .numpy import TreeNumpy
 from ..common import ireduce, TreeObject
-from ..utils import replaceable_partial
+from ..utils import replaceable_partial, inherit_doc
 
 __all__ = [
     'all', 'any',
     'equal', 'array_equal',
 ]
 
+
+def _doc_stripper(src, _, lines: List[str]):
+    _name, _version = src.__name__, np.__version__
+    _short_version = '.'.join(_version.split('.')[:2])
+    return [
+        f'.. note::',
+        f'',
+        f'    This documentation is based on '
+        f'    `numpy.{_name} <https://numpy.org/doc/{_short_version}/reference/generated/numpy.{_name}.html>`_ '
+        f'    in `numpy v{_version} <https://numpy.org/doc/{_short_version}/>`_.',
+        f'    **Its arguments\' arrangements depend on the version of numpy you installed**.',
+        f'',
+        *lines,
+    ]
+
+
 func_treelize = replaceable_partial(original_func_treelize, return_type=TreeNumpy)
+docs = replaceable_partial(inherit_doc, stripper=_doc_stripper)
 
 
+@docs(np.all)
 @ireduce(builtins.all)
 @func_treelize(return_type=TreeObject)
 def all(a, *args, **kwargs):
     return np.all(a, *args, **kwargs)
 
 
+@docs(np.any)
 @ireduce(builtins.any)
 @func_treelize()
 def any(a, *args, **kwargs):
     return np.any(a, *args, **kwargs)
 
 
+@docs(np.equal)
 @func_treelize()
 def equal(x1, x2, *args, **kwargs):
     return np.equal(x1, x2, *args, **kwargs)
 
 
+@docs(np.array_equal)
 @func_treelize()
 def array_equal(a1, a2, *args, **kwargs):
     return np.array_equal(a1, a2, *args, **kwargs)

treetensor/torch/funcs.py

@@ -9,6 +9,17 @@ from .tensor import TreeTensor, tireduce
 from ..common import TreeObject, ireduce
 from ..utils import replaceable_partial, direct_doc, inherit_doc
 
+__all__ = [
+    'zeros', 'zeros_like',
+    'randn', 'randn_like',
+    'randint', 'randint_like',
+    'ones', 'ones_like',
+    'full', 'full_like',
+    'empty', 'empty_like',
+    'all', 'any',
+    'eq', 'equal',
+]
+
 
 def _doc_stripper(src, _, lines: List[str]):
     _name, _version = src.__name__, torch.__version__
@@ -30,17 +41,6 @@ def _doc_stripper(src, _, lines: List[str]):
 func_treelize = replaceable_partial(original_func_treelize, return_type=TreeTensor)
 docs = post_process(post_process(direct_doc))(replaceable_partial(inherit_doc, stripper=_doc_stripper))
 
-__all__ = [
-    'zeros', 'zeros_like',
-    'randn', 'randn_like',
-    'randint', 'randint_like',
-    'ones', 'ones_like',
-    'full', 'full_like',
-    'empty', 'empty_like',
-    'all', 'any',
-    'eq', 'equal',
-]
-
 
 @docs(torch.zeros)
 @func_treelize()

treetensor/utils/doc.py

@@ -7,7 +7,9 @@ __all__ = [
 
 
 def _strip_lines(doc: Optional[str]):
-    _lines = (doc or '').strip().splitlines()
+    _lines = (doc or '').splitlines()
+    _first_line_id, _ = sorted(filter(lambda t: t[1].strip(), enumerate(_lines)))[0]
+    _lines = _lines[_first_line_id:]
     _exist_lines = list(filter(str.strip, _lines))
 
     if not _exist_lines: