Merge pull request #80 from tensorcircuit/perf-fix-arg-alias-18006180529273122311

⚡ Optimize arg_alias and fix string alias bug

Author: refraction-ray

tensorcircuit/utils.py

@@ -149,16 +149,24 @@ def arg_alias(
     :rtype: Callable[..., Any]
     """
 
+    flat_aliases = []
+    normalized_alias_dict = {}
+    for k, vs in alias_dict.items():
+        if isinstance(vs, str):
+            vs_list = [vs]
+        else:
+            vs_list = vs  # type: ignore
+        normalized_alias_dict[k] = vs_list
+        for v in vs_list:
+            flat_aliases.append((k, v))
+
     @wraps(f)
     def wrapper(*args: Any, **kws: Any) -> Any:
-        for k, vs in alias_dict.items():
-            if isinstance(vs, str):
-                vs = []
-            for v in vs:
-                if v in kws:
-                    # in case it is None by design!
-                    kws[k] = kws[v]
-                    del kws[v]
+        for k, v in flat_aliases:
+            if v in kws:
+                # in case it is None by design!
+                kws[k] = kws[v]
+                del kws[v]
         return f(*args, **kws)
 
     if fix_doc:
@@ -176,15 +184,15 @@ def wrapper(*args: Any, **kws: Any) -> Any:
 
                 if line.strip().startswith(":param "):
                     param = re.findall(r":param\s([^\s]+):", line)[0]
-                    if param in alias_dict:
+                    if param in normalized_alias_dict:
                         j = 1
                         while True:
                             ndoc.append(doc[i + j])
                             if doc[i + j].strip().startswith(":type"):
                                 break
                             j += 1
                         skip = True
-                        for v in alias_dict[param]:
+                        for v in normalized_alias_dict[param]:
                             ndoc.append(
                                 re.sub(
                                     r"(.*:param\s)([^\s]+)(:.*)",

tests/test_miscs.py

@@ -192,6 +192,20 @@ def f(theta: float, beta: float) -> float:
     print(f.__doc__)
     assert len(f.__doc__.strip().split("\n")) == 12
 
+    @partial(tc.utils.arg_alias, alias_dict={"theta": "alpha"})
+    def g(theta: float) -> float:
+        """
+        g doc
+
+        :param theta: theta angle
+        :type theta: float
+        :return: theta
+        """
+        return theta
+
+    np.testing.assert_allclose(g(alpha=1.0), 1.0, atol=1e-5)
+    assert "alpha: alias for the argument ``theta``" in g.__doc__
+
 
 def test_finite_difference_tf(tfb):
     def f(param1, param2):