tests: more fun with workers making graphs

justfile

@@ -59,6 +59,7 @@ generate:
   just stubs-generate-new 'tierkreis_workers/pytket_worker'
   just stubs-generate 'tierkreis_workers/quantinuum_worker'
   just stubs-generate 'tierkreis_workers/qulacs_worker'
+  just stubs-generate 'tierkreis/tests/workers/graph'
 
   just stubs-generate-api 'docs/source/examples/example_workers/auth_worker'
   just stubs-generate-api 'docs/source/examples/example_workers/error_worker'

tierkreis/tests/controller/test_resume.py

@@ -24,6 +24,8 @@
 )
 from tests.controller.typed_graphdata import (
     eval_body_is_from_worker,
+    eval_graph_of_graph,
+    eval_from_worker_with_graph_from_worker,
     factorial,
     gcd,
     tkr_conj,
@@ -214,14 +216,26 @@ def test_resume(  # noqa: PLR0913
     tuple[GraphData | GraphBuilder, Any, str, dict[str, PType] | PType]
 ] = [
     (eval_body_is_from_worker(), 21, "eval_body_is_from_worker", {"value": 10}),
+    (eval_graph_of_graph(), 31, "eval_graph_of_graph", {"value": 3}),
+    (
+        eval_from_worker_with_graph_from_worker(),
+        23,
+        "eval_from_worker_with_graph_from_worker",
+        {"value": 5},
+    ),
 ]
+
 with_worker_params: list[
     tuple[GraphData | GraphBuilder, Any, str, int, dict[str, PType] | PType]
 ] = [
     (graph, output, name, i + 1, inputs)
     for i, (graph, output, name, inputs) in enumerate(with_worker_param_data)
 ]
-with_worker_ids = ["eval_body_is_from_worker"]
+with_worker_ids = [
+    "eval_body_is_from_worker",
+    "eval_graph_of_graph",
+    "eval_from_worker_with_graph_from_worker",
+]
 
 
 @pytest.mark.parametrize(

tierkreis/tests/controller/typed_graphdata.py

@@ -1,6 +1,6 @@
 from typing import NamedTuple
 
-from tests.workers.graph.stubs import doubler_plus_graph
+from tests.workers.graph.stubs import doubler_plus_graph, graph_of_graph, apply_twice
 from tierkreis.builder import GraphBuilder, TypedGraphRef
 from tierkreis.builtins import (
     conjugate,
@@ -15,7 +15,7 @@
     tkr_str,
 )
 from tierkreis.controller.data.core import EmptyModel
-from tierkreis.controller.data.models import TKR
+from tierkreis.controller.data.models import TKR, OpaqueType
 
 
 class DoublerInput(NamedTuple):
@@ -167,6 +167,40 @@ def eval_body_is_from_worker() -> GraphBuilder[TKR[int], TKR[int]]:
     return g
 
 
+class ApplyTwiceInput(NamedTuple):
+    # Note we mangle this as the stub generator does, unlike
+    # the original worker ApplyTwiceInput which takes TKR[GraphData]
+    graph: TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]  # noqa: F821
+    value: TKR[int]
+
+
+def eval_from_worker_with_graph_from_worker() -> GraphBuilder[TKR[int], TKR[int]]:
+    g = GraphBuilder(TKR[int], TKR[int])
+    graph = g.task(doubler_plus_graph())
+    # This is ok, but we can't pass the graph_ref into ApplyTwiceInputs
+    # graph_ref = TypedGraphRef(graph.value_ref(), TKR[int], TKR[int])
+    inputs = ApplyTwiceInput(graph=graph, value=g.inputs)
+
+    ap2 = g.task(apply_twice())
+    # TypedGraphRef Constructor takes outputs then inputs, but type-params are inputs then outputs:
+    ap2_ref = TypedGraphRef(ap2.value_ref(), TKR[int], ApplyTwiceInput)
+    out = g.eval(ap2_ref, inputs)
+    g.outputs(out)
+    return g
+
+
+def eval_graph_of_graph() -> GraphBuilder[TKR[int], TKR[int]]:
+    g = GraphBuilder(TKR[int], TKR[int])
+    graph = g.task(doubler_plus_graph())
+    # This is ok, but we can't pass the graph_ref into exponentiate_graph:
+    # graph_ref = TypedGraphRef(graph.value_ref(), TKR[int], TKR[int])
+    eg = g.task(graph_of_graph(graph, g.const(3)))
+    exp_graph = TypedGraphRef(eg.value_ref(), TKR[int], TKR[int])
+    out = g.eval(exp_graph, g.inputs)
+    g.outputs(out)
+    return g
+
+
 def embed_graph():
     class InnerOutput(NamedTuple):
         log: TKR[str]

tierkreis/tests/workers/graph/main.py

@@ -1,7 +1,8 @@
 from sys import argv
+from typing import NamedTuple
 
 from tierkreis import Worker
-from tierkreis.builder import GraphBuilder
+from tierkreis.builder import GraphBuilder, TypedGraphRef
 from tierkreis.builtins import iadd, itimes
 from tierkreis.controller.data.graph import GraphData
 from tierkreis.models import TKR
@@ -18,5 +19,41 @@ def doubler_plus_graph() -> GraphData:
     return g.get_data()
 
 
+@worker.task()
+# The input graph here is expected to be int->int, but we have no way to express that in the type system.
+# (GraphBuilder doesn't work as it's not accepted by the stub generator)
+def graph_of_graph(f: GraphData, n: int) -> GraphData:
+    """Builds a new graph: lambda x: f^n(x)
+
+    I.e. the graph applies the first argument `f` to the graph's input `n` times.
+    The graph contains the argument graph `f` as a constant."""
+    g = GraphBuilder(TKR[int], TKR[int])
+    v = g.inputs
+    ref = TypedGraphRef(g.const(f).value_ref(), TKR[int], TKR[int])
+    for _ in range(n):
+        v = g.eval(ref, v)
+    g.outputs(v)
+    return g.get_data()
+
+
+class ApplyTwiceInput(NamedTuple):
+    graph: TKR[GraphData]
+    value: TKR[int]
+
+
+@worker.task()
+def apply_twice() -> GraphData:
+    """Returns a graph for lambda f,x: f(f(x)).
+
+    That is, `f` and `x` are inputs to the graph, not the worker function building it.
+    """
+    g = GraphBuilder(ApplyTwiceInput, TKR[int])
+    f = TypedGraphRef(g.inputs.graph.value_ref(), TKR[int], TKR[int])
+    run_once = g.eval(f, g.inputs.value)
+    run_twice = g.eval(f, run_once)
+    g.outputs(run_twice)
+    return g.get_data()
+
+
 if __name__ == "__main__":
     worker.app(argv)

tierkreis/tests/workers/graph/stubs.py

@@ -12,3 +12,26 @@ def out() -> type[TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]]:
     @property
     def namespace(self) -> str:
         return "graph"
+
+
+class graph_of_graph(NamedTuple):
+    f: TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]  # noqa: F821 # fmt: skip
+    n: TKR[int]  # noqa: F821 # fmt: skip
+
+    @staticmethod
+    def out() -> type[TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]]:  # noqa: F821 # fmt: skip
+        return TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]  # noqa: F821 # fmt: skip
+
+    @property
+    def namespace(self) -> str:
+        return "graph"
+
+
+class apply_twice(NamedTuple):
+    @staticmethod
+    def out() -> type[TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]]:  # noqa: F821 # fmt: skip
+        return TKR[OpaqueType["tierkreis.controller.data.graph.GraphData"]]  # noqa: F821 # fmt: skip
+
+    @property
+    def namespace(self) -> str:
+        return "graph"