diff --git a/justfile b/justfile index ae5b2bb22..1bc381311 100644 --- a/justfile +++ b/justfile @@ -48,6 +48,8 @@ stubs-generate-api dir: generate: just stubs-generate 'tierkreis/tierkreis/builtins' + just stubs-generate 'tierkreis/tests/workers/graph' + just stubs-generate-api 'tierkreis_workers/aer_worker/tkr_aer_worker_impl' just stubs-generate-api 'tierkreis_workers/ibmq_worker/tkr_ibmq_worker_impl' just stubs-generate-api 'tierkreis_workers/nexus_worker/tkr_nexus_worker_impl' diff --git a/tierkreis/tests/controller/test_resume.py b/tierkreis/tests/controller/test_resume.py index 42efc6ca2..af04fe824 100644 --- a/tierkreis/tests/controller/test_resume.py +++ b/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( diff --git a/tierkreis/tests/controller/typed_graphdata.py b/tierkreis/tests/controller/typed_graphdata.py index c4c787afd..669fadad3 100644 --- a/tierkreis/tests/controller/typed_graphdata.py +++ b/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,39 @@ def eval_body_is_from_worker() -> GraphBuilder[TKR[int], TKR[int]]: return g +class ApplyTwiceInput(NamedTuple): + # Note we mangle this like the stub generator would, although the generator + # never sees the graph's inputs as they are hidden in an untyped 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 ApplyTwiceInput + # graph_ref = TypedGraphRef(graph.value_ref(), TKR[int], TKR[int]) + inputs = ApplyTwiceInput(graph=graph, value=g.inputs) + + ap2 = g.task(apply_twice()) + ap2_ref = TypedGraphRef(ap2.value_ref(), ApplyTwiceInput, TKR[int]) + 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] diff --git a/tierkreis/tests/workers/graph/main.py b/tierkreis/tests/workers/graph/main.py index 7b21983dc..65154cfe1 100644 --- a/tierkreis/tests/workers/graph/main.py +++ b/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() + + +@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. + """ + + class ApplyTwiceInput(NamedTuple): + graph: TKR[GraphData] + value: TKR[int] + + 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) diff --git a/tierkreis/tests/workers/graph/stubs.py b/tierkreis/tests/workers/graph/stubs.py index 5610d2e03..9217da5c0 100644 --- a/tierkreis/tests/workers/graph/stubs.py +++ b/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"