add node capture infrastructure

Author: refraction-ray

CHANGELOG.md

@@ -12,6 +12,8 @@
 
 - Add `with_prob` for `stabilizercircuit.measure()`.
 
+- Add `tc.cons.function_nodes_capture` decorator and `tc.cons.runtime_nodes_capture` context manager for directly return nodes before real contraction.
+
 ### Fixed
 
 - Fix the nodes order in contraction by giving each node a global `_stable_id_`.

tensorcircuit/cons.py

@@ -516,8 +516,8 @@ def _get_path_cache_friendly(
     nodes = list(nodes)
 
     nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
-    if isinstance(algorithm, list):
-        return algorithm, nodes_new
+    # if isinstance(algorithm, list):
+    #     return algorithm, [nodes_new]
 
     all_edges = tn.get_all_edges(nodes_new)
     all_edges_sorted = sorted_edges(all_edges)
@@ -693,6 +693,51 @@ def _base(
     return final_node
 
 
+class NodesReturn(Exception):
+    """
+    Intentionally stop execution to return a value.
+    """
+
+    def __init__(self, value_to_return: Any):
+        self.value = value_to_return
+        super().__init__(
+            f"Intentionally stopping execution to return: {value_to_return}"
+        )
+
+
+def _get_sorted_nodes(nodes: List[Any], *args: Any, **kws: Any) -> Any:
+    nodes_new = sorted(nodes, key=lambda node: getattr(node, "_stable_id_", -1))
+    raise NodesReturn(nodes_new)
+
+
+def function_nodes_capture(func: Callable[[Any], Any]) -> Callable[[Any], Any]:
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> Any:
+        with runtime_contractor(method="before"):
+            try:
+                result = func(*args, **kwargs)
+                return result
+            except NodesReturn as e:
+                return e.value
+
+    return wrapper
+
+
+@contextmanager
+def runtime_nodes_capture(key: str = "nodes") -> Iterator[Any]:
+    old_contractor = getattr(thismodule, "contractor")
+    set_contractor(method="before")
+    captured_value: Dict[str, List[tn.Node]] = {}
+    try:
+        yield captured_value
+    except NodesReturn as e:
+        captured_value[key] = e.value
+    finally:
+        for module in sys.modules:
+            if module.startswith(package_name):
+                setattr(sys.modules[module], "contractor", old_contractor)
+
+
 def custom(
     nodes: List[Any],
     optimizer: Any,
@@ -763,6 +808,16 @@ def custom_stateful(
 
 # only work for custom
 def contraction_info_decorator(algorithm: Callable[..., Any]) -> Callable[..., Any]:
+    """Decorator to add contraction information logging to an optimizer.
+
+    This decorator wraps an optimization algorithm and prints detailed information
+    about the contraction cost (FLOPs, size, write) and path finding time.
+
+    :param algorithm: The optimization algorithm to decorate.
+    :type algorithm: Callable[..., Any]
+    :return: The decorated optimization algorithm.
+    :rtype: Callable[..., Any]
+    """
     from cotengra import ContractionTree
 
     def new_algorithm(
@@ -869,6 +924,9 @@ def set_contractor(
             **kws,
         )
 
+    elif method == "before":  # a hack way to get the nodes
+        cf = _get_sorted_nodes
+
     else:
         # cf = getattr(tn.contractors, method, None)
         # if not cf:

tests/test_circuit.py

@@ -889,6 +889,19 @@ def test_circuit_quoperator(backend):
     np.testing.assert_allclose(qo.eval_matrix(), c.matrix(), atol=1e-5)
 
 
+def test_perm_matrix():
+    from tensorcircuit.translation import perm_matrix
+
+    p2 = perm_matrix(2)
+    np.testing.assert_allclose(
+        p2, np.array([[1, 0, 0, 0], [0, 0, 1, 0], [0, 1, 0, 0], [0, 0, 0, 1]])
+    )
+    p3 = perm_matrix(3)
+    v = np.arange(8)
+    vt = np.array([0, 4, 2, 6, 1, 5, 3, 7])
+    np.testing.assert_allclose(p3 @ v, vt)
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
 def test_qir2cirq(backend):
     try:

tests/test_miscs.py

@@ -312,3 +312,23 @@ def baseline(params):
         return c.expectation_ps(z=[-1])
 
     np.testing.assert_allclose(value, baseline(params), atol=1e-6)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_runtime_nodes_capture(backend):
+    with tc.cons.runtime_nodes_capture() as captured:
+        c = tc.Circuit(3)
+        c.h(0)
+        c.amplitude("010")
+    len(captured["nodes"]) == 7
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_function_nodes_capture(backend):
+    @tc.cons.function_nodes_capture
+    def exp(theta):
+        c = tc.Circuit(3)
+        c.h(0)
+        return c.expectation_ps(z=[-3], reuse=False)
+
+    assert len(exp(0.3)) == 9