Repeat boxes working.

examples/Model definition

@@ -438,8 +438,8 @@
       "height": 200.0,
       "id": "MSE loss 2",
       "position": {
-        "x": 690.0,
-        "y": -480.0
+        "x": 309.4422414664647,
+        "y": -552.1056805642488
       },
       "type": "basic",
       "width": 200.0

examples/Model use

@@ -579,54 +579,54 @@
               ],
               "data": [
                 [
-                  "[0.33480108 0.59181517 0.76198453 0.98062384]",
-                  "[1.33480108 1.59181523 1.76198459 1.98062384]",
-                  "[1.3419755697250366, 1.5946478843688965, 1.7717586755752563, 1.9897401332855225]"
+                  "[0.94559073 0.65736622 0.25761551 0.48553199]",
+                  "[1.94559073 1.65736628 1.25761557 1.48553205]",
+                  "[1.5948047637939453, 1.619612693786621, 1.5269112586975098, -0.008817584253847599]"
                 ],
                 [
-                  "[0.91730917 0.22574073 0.09591609 0.33056474]",
-                  "[1.91730917 1.22574067 1.09591603 1.33056474]",
-                  "[1.900892972946167, 1.2247941493988037, 1.0862866640090942, 1.323314905166626]"
+                  "[0.47856545 0.46267092 0.6376707  0.84747767]",
+                  "[1.47856545 1.46267092 1.63767076 1.84747767]",
+                  "[1.5928349494934082, 1.6176562309265137, 1.52553391456604, -0.008808750659227371]"
                 ],
                 [
-                  "[0.32565445 0.90939188 0.07488042 0.13730896]",
-                  "[1.32565451 1.90939188 1.07488036 1.13730896]",
-                  "[1.3460955619812012, 1.8960161209106445, 1.0530263185501099, 1.1075329780578613]"
+                  "[0.59492421 0.90274489 0.38069052 0.46101224]",
+                  "[1.59492421 1.90274489 1.38069057 1.46101224]",
+                  "[1.5592631101608276, 1.5841729640960693, 1.5020174980163574, -0.008660631254315376]"
                 ],
                 [
-                  "[0.87608397 0.93200487 0.80169648 0.37758952]",
-                  "[1.87608397 1.93200493 1.80169654 1.37758946]",
-                  "[1.87070894241333, 1.9386992454528809, 1.8151044845581055, 1.3952441215515137]"
+                  "[0.12858278 0.09930819 0.83222693 0.72485673]",
+                  "[1.12858272 1.09930825 1.83222699 1.72485673]",
+                  "[1.469609260559082, 1.4953691959381104, 1.4395854473114014, -0.00825517252087593]"
                 ],
                 [
-                  "[0.39147133 0.29854035 0.84663737 0.58175623]",
-                  "[1.39147139 1.29854035 1.84663737 1.58175623]",
-                  "[1.3877646923065186, 1.2995290756225586, 1.847062587738037, 1.583693265914917]"
+                  "[0.94516498 0.08422136 0.5608117  0.07652664]",
+                  "[1.94516492 1.08422136 1.56081176 1.07652664]",
+                  "[1.5648787021636963, 1.5899150371551514, 1.5060429573059082, -0.008683123625814915]"
                 ],
                 [
-                  "[0.48507756 0.80808765 0.77162558 0.47834778]",
-                  "[1.48507762 1.80808759 1.77162552 1.47834778]",
-                  "[1.490919828414917, 1.8087174892425537, 1.7757861614227295, 1.4824031591415405]"
+                  "[0.54914117 0.03810108 0.87531954 0.73044223]",
+                  "[1.54914117 1.03810108 1.87531948 1.73044229]",
+                  "[1.6262837648391724, 1.650843620300293, 1.548863410949707, -0.00895910244435072]"
                 ],
                 [
-                  "[0.75292218 0.81470108 0.49657214 0.56217098]",
-                  "[1.75292218 1.81470108 1.49657214 1.56217098]",
-                  "[1.7527031898498535, 1.8176040649414062, 1.503413438796997, 1.570152759552002]"
+                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
+                  "[1.94221354 1.57740951 1.98649526 1.40934443]",
+                  "[1.8493703603744507, 1.8721930980682373, 1.704444169998169, -0.009961890056729317]"
                 ],
                 [
-                  "[0.11693293 0.49860179 0.55020827 0.88832849]",
-                  "[1.11693287 1.49860179 1.55020833 1.88832855]",
-                  "[1.1314976215362549, 1.4944026470184326, 1.546830177307129, 1.8803892135620117]"
+                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
+                  "[1.8065424  1.08253479 1.74478531 1.71257162]",
+                  "[1.672502040863037, 1.6967051029205322, 1.5810983180999756, -0.009166811592876911]"
                 ],
                 [
-                  "[0.19409031 0.68692201 0.60667384 0.57829887]",
-                  "[1.19409037 1.68692207 1.60667384 1.57829881]",
-                  "[1.2091591358184814, 1.6816589832305908, 1.6011345386505127, 1.5684995651245117]"
+                  "[0.50272274 0.54912758 0.17663097 0.79070699]",
+                  "[1.50272274 1.54912758 1.17663097 1.79070699]",
+                  "[1.4309396743774414, 1.4570224285125732, 1.4126286506652832, -0.008081027306616306]"
                 ],
                 [
-                  "[0.62569475 0.9881897  0.83639616 0.9828859 ]",
-                  "[1.62569475 1.9881897  1.83639622 1.98288584]",
-                  "[1.6314740180969238, 1.996805191040039, 1.8592857122421265, 2.0075552463531494]"
+                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
+                  "[1.34084332 1.73018837 1.54168713 1.91440833]",
+                  "[1.5581963062286377, 1.5832865238189697, 1.5013742446899414, -0.008653069846332073]"
                 ]
               ]
             },
@@ -648,6 +648,10 @@
                   "[0.11560339 0.57495481 0.76535827 0.0391947 ]",
                   "[1.11560345 1.57495475 1.76535821 1.0391947 ]"
                 ],
+                [
+                  "[0.19409031 0.68692201 0.60667384 0.57829887]",
+                  "[1.19409037 1.68692207 1.60667384 1.57829881]"
+                ],
                 [
                   "[0.76807946 0.98855817 0.08259124 0.01730657]",
                   "[1.76807952 1.98855817 1.0825913  1.01730657]"
@@ -681,12 +685,12 @@
                   "[1.98324287 1.99464178 1.14008355 1.47651017]"
                 ],
                 [
-                  "[0.48959708 0.48549271 0.32688856 0.356677  ]",
-                  "[1.48959708 1.48549271 1.32688856 1.35667706]"
+                  "[0.11693293 0.49860179 0.55020827 0.88832849]",
+                  "[1.11693287 1.49860179 1.55020833 1.88832855]"
                 ],
                 [
-                  "[0.50272274 0.54912758 0.17663097 0.79070699]",
-                  "[1.50272274 1.54912758 1.17663097 1.79070699]"
+                  "[0.48959708 0.48549271 0.32688856 0.356677  ]",
+                  "[1.48959708 1.48549271 1.32688856 1.35667706]"
                 ],
                 [
                   "[0.04508126 0.76880038 0.80721325 0.62542385]",
@@ -708,6 +712,10 @@
                   "[0.24388778 0.07268471 0.68350857 0.73431659]",
                   "[1.24388778 1.07268476 1.68350863 1.73431659]"
                 ],
+                [
+                  "[0.62569475 0.9881897  0.83639616 0.9828859 ]",
+                  "[1.62569475 1.9881897  1.83639622 1.98288584]"
+                ],
                 [
                   "[0.56922203 0.98222166 0.76851749 0.28615737]",
                   "[1.56922197 1.9822216  1.76851749 1.28615737]"
@@ -720,6 +728,10 @@
                   "[0.90817457 0.89270043 0.38583666 0.66566533]",
                   "[1.90817451 1.89270043 1.3858366  1.66566539]"
                 ],
+                [
+                  "[0.48507756 0.80808765 0.77162558 0.47834778]",
+                  "[1.48507762 1.80808759 1.77162552 1.47834778]"
+                ],
                 [
                   "[0.68062544 0.98093534 0.14778823 0.53244978]",
                   "[1.68062544 1.98093534 1.14778829 1.53244972]"
@@ -740,10 +752,6 @@
                   "[0.23942459 0.90487361 0.69337189 0.65089428]",
                   "[1.23942459 1.90487361 1.69337189 1.65089428]"
                 ],
-                [
-                  "[0.94516498 0.08422136 0.5608117  0.07652664]",
-                  "[1.94516492 1.08422136 1.56081176 1.07652664]"
-                ],
                 [
                   "[0.26661873 0.45946234 0.13510543 0.81294441]",
                   "[1.26661873 1.4594624  1.13510537 1.81294441]"
@@ -772,10 +780,6 @@
                   "[0.78956431 0.87284744 0.06880784 0.03455889]",
                   "[1.78956437 1.87284744 1.06880784 1.03455889]"
                 ],
-                [
-                  "[0.94221359 0.57740951 0.98649532 0.40934443]",
-                  "[1.94221354 1.57740951 1.98649526 1.40934443]"
-                ],
                 [
                   "[0.00497234 0.39319336 0.57054168 0.75150961]",
                   "[1.00497234 1.39319336 1.57054162 1.75150967]"
@@ -788,6 +792,14 @@
                   "[0.72290605 0.96945059 0.68354797 0.15270454]",
                   "[1.72290611 1.96945059 1.68354797 1.15270448]"
                 ],
+                [
+                  "[0.75292218 0.81470108 0.49657214 0.56217098]",
+                  "[1.75292218 1.81470108 1.49657214 1.56217098]"
+                ],
+                [
+                  "[0.33480108 0.59181517 0.76198453 0.98062384]",
+                  "[1.33480108 1.59181523 1.76198459 1.98062384]"
+                ],
                 [
                   "[0.52784437 0.54268694 0.12358981 0.72116476]",
                   "[1.52784443 1.54268694 1.12358975 1.7211647 ]"
@@ -796,10 +808,6 @@
                   "[0.73217702 0.65233225 0.44077861 0.33837909]",
                   "[1.73217702 1.65233231 1.44077861 1.33837914]"
                 ],
-                [
-                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
-                  "[1.34084332 1.73018837 1.54168713 1.91440833]"
-                ],
                 [
                   "[0.60110539 0.3618983  0.32342511 0.98672163]",
                   "[1.60110545 1.3618983  1.32342505 1.98672163]"
@@ -836,6 +844,10 @@
                   "[0.95928186 0.84273899 0.71514636 0.38619852]",
                   "[1.95928192 1.84273899 1.7151463  1.38619852]"
                 ],
+                [
+                  "[0.32565445 0.90939188 0.07488042 0.13730896]",
+                  "[1.32565451 1.90939188 1.07488036 1.13730896]"
+                ],
                 [
                   "[0.9829582  0.59269661 0.40120947 0.95487177]",
                   "[1.9829582  1.59269667 1.40120947 1.95487177]"
@@ -844,10 +856,6 @@
                   "[0.79905868 0.89367443 0.75429088 0.3190186 ]",
                   "[1.79905868 1.89367437 1.75429082 1.3190186 ]"
                 ],
-                [
-                  "[0.54914117 0.03810108 0.87531954 0.73044223]",
-                  "[1.54914117 1.03810108 1.87531948 1.73044229]"
-                ],
                 [
                   "[0.67418337 0.79634351 0.23229051 0.71345252]",
                   "[1.67418337 1.79634356 1.23229051 1.71345258]"
@@ -860,14 +868,14 @@
                   "[0.81788456 0.58174163 0.29376316 0.7971254 ]",
                   "[1.81788456 1.58174157 1.29376316 1.79712534]"
                 ],
-                [
-                  "[0.94559073 0.65736622 0.25761551 0.48553199]",
-                  "[1.94559073 1.65736628 1.25761557 1.48553205]"
-                ],
                 [
                   "[0.60075855 0.12234765 0.00614399 0.30560958]",
                   "[1.60075855 1.12234759 1.00614405 1.30560958]"
                 ],
+                [
+                  "[0.39147133 0.29854035 0.84663737 0.58175623]",
+                  "[1.39147139 1.29854035 1.84663737 1.58175623]"
+                ],
                 [
                   "[0.02162331 0.81861657 0.92468154 0.07808572]",
                   "[1.02162337 1.81861663 1.92468154 1.07808566]"
@@ -896,10 +904,6 @@
                   "[0.60609657 0.96257663 0.19292736 0.95702219]",
                   "[1.60609651 1.96257663 1.19292736 1.95702219]"
                 ],
-                [
-                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
-                  "[1.8065424  1.08253479 1.74478531 1.71257162]"
-                ],
                 [
                   "[0.70167565 0.26930219 0.5660674  0.61194974]",
                   "[1.70167565 1.26930213 1.56606746 1.61194968]"
@@ -908,10 +912,6 @@
                   "[0.76933283 0.86241865 0.44114518 0.65644735]",
                   "[1.76933289 1.86241865 1.44114518 1.65644741]"
                 ],
-                [
-                  "[0.59492421 0.90274489 0.38069052 0.46101224]",
-                  "[1.59492421 1.90274489 1.38069057 1.46101224]"
-                ],
                 [
                   "[0.15064228 0.03198934 0.25754827 0.51484001]",
                   "[1.15064228 1.03198934 1.25754833 1.51484001]"
@@ -920,6 +920,10 @@
                   "[0.12024075 0.21342516 0.56858408 0.58644271]",
                   "[1.12024069 1.21342516 1.56858408 1.58644271]"
                 ],
+                [
+                  "[0.91730917 0.22574073 0.09591609 0.33056474]",
+                  "[1.91730917 1.22574067 1.09591603 1.33056474]"
+                ],
                 [
                   "[0.49691743 0.61873293 0.90698647 0.94486356]",
                   "[1.49691749 1.61873293 1.90698647 1.94486356]"
@@ -948,18 +952,10 @@
                   "[0.80893755 0.92237449 0.88346356 0.93164903]",
                   "[1.80893755 1.92237449 1.88346362 1.93164897]"
                 ],
-                [
-                  "[0.12858278 0.09930819 0.83222693 0.72485673]",
-                  "[1.12858272 1.09930825 1.83222699 1.72485673]"
-                ],
                 [
                   "[0.72470158 0.4940322  0.41027349 0.89364016]",
                   "[1.72470164 1.49403214 1.41027355 1.89364016]"
                 ],
-                [
-                  "[0.47856545 0.46267092 0.6376707  0.84747767]",
-                  "[1.47856545 1.46267092 1.63767076 1.84747767]"
-                ],
                 [
                   "[0.49584109 0.80599248 0.07096875 0.75872749]",
                   "[1.49584103 1.80599248 1.07096875 1.75872755]"
@@ -992,6 +988,10 @@
                   "[0.72795159 0.79317838 0.27832931 0.96576637]",
                   "[1.72795153 1.79317832 1.27832937 1.96576643]"
                 ],
+                [
+                  "[0.87608397 0.93200487 0.80169648 0.37758952]",
+                  "[1.87608397 1.93200493 1.80169654 1.37758946]"
+                ],
                 [
                   "[0.68891573 0.25576538 0.96339929 0.503833  ]",
                   "[1.68891573 1.25576544 1.96339929 1.50383306]"
@@ -1000,7 +1000,7 @@
             }
           },
           "other": {
-            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__tensor_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (2) - <function leaky_relu at 0x759ed4f2c360>: Linear_2_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (4) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__tensor_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['END_Repeat_1_output', 'Input__tensor_3_x'], loss=Sequential(\n  (0) - <function mse_loss at 0x759ed4f2de40>: END_Repeat_1_output, Input__tensor_3_x -> MSE_loss_2_output\n  (1) - Identity(): MSE_loss_2_output -> loss\n), optimizer=SGD (\nParameter Group 0\n    dampening: 0\n    differentiable: False\n    foreach: None\n    fused: None\n    lr: 0.1\n    maximize: False\n    momentum: 0\n    nesterov: False\n    weight_decay: 0\n), source_workspace=None, trained=True)"
+            "model": "ModelConfig(model=Sequential(\n  (0) - Identity(): Input__tensor_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (2) - <function leaky_relu at 0x75e660938220>: Linear_2_output -> Activation_1_output\n  (3) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (4) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (5) - <function leaky_relu at 0x75e660938220>: Linear_2_output -> Activation_1_output\n  (6) - Identity(): Activation_1_output -> START_Repeat_1_output\n  (7) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_2_output\n  (8) - <function leaky_relu at 0x75e660938220>: Linear_2_output -> Activation_1_output\n  (9) - Identity(): Activation_1_output -> END_Repeat_1_output\n  (10) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__tensor_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['Input__tensor_3_x', 'END_Repeat_1_output'], loss=Sequential(\n  (0) - <function mse_loss at 0x75e660939d00>: END_Repeat_1_output, Input__tensor_3_x -> MSE_loss_2_output\n  (1) - Identity(): MSE_loss_2_output -> loss\n), optimizer=SGD (\nParameter Group 0\n    dampening: 0\n    differentiable: False\n    foreach: None\n    fused: None\n    lr: 0.1\n    maximize: False\n    momentum: 0\n    nesterov: False\n    weight_decay: 0\n), source_workspace=None, trained=True)"
           },
           "relations": []
         },
@@ -1035,8 +1035,8 @@
                     "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "END_Repeat_1_output",
-                    "Input__tensor_3_x"
+                    "Input__tensor_3_x",
+                    "END_Repeat_1_output"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1210,8 +1210,8 @@
                     "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "END_Repeat_1_output",
-                    "Input__tensor_3_x"
+                    "Input__tensor_3_x",
+                    "END_Repeat_1_output"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"
@@ -1270,7 +1270,7 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "1500",
+          "epochs": "110",
           "input_mapping": "{\"map\":{\"Input__tensor_1_x\":{\"df\":\"df_train\",\"column\":\"x\"},\"Input__tensor_3_x\":{\"df\":\"df_train\",\"column\":\"y\"}}}",
           "model_name": "model"
         },
@@ -1322,8 +1322,8 @@
                     "Input__tensor_1_x"
                   ],
                   "loss_inputs": [
-                    "END_Repeat_1_output",
-                    "Input__tensor_3_x"
+                    "Input__tensor_3_x",
+                    "END_Repeat_1_output"
                   ],
                   "outputs": [
                     "END_Repeat_1_output"

lynxkite-core/src/lynxkite/core/ops.py

@@ -13,7 +13,7 @@ from typing_extensions import Annotated
 if typing.TYPE_CHECKING:
     from . import workspace
 
-CATALOGS = {}
+CATALOGS: dict[str, dict[str, "Op"]] = {}
 EXECUTORS = {}
 
 typeof = type  # We have some arguments called "type".

lynxkite-graph-analytics/src/lynxkite_graph_analytics/pytorch_model_ops.py

@@ -214,6 +214,9 @@ class ModelConfig:
     source_workspace: str | None = None
     trained: bool = False
 
+    def num_parameters(self) -> int:
+        return sum(p.numel() for p in self.model.parameters())
+
     def _forward(self, inputs: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
         model_inputs = [inputs[i] for i in self.model_inputs]
         output = self.model(*model_inputs)
@@ -270,7 +273,7 @@ class ModelBuilder:
     def __init__(self, ws: workspace.Workspace, inputs: dict[str, torch.Tensor]):
         self.catalog = ops.CATALOGS[ENV]
         optimizers = []
-        self.nodes = {}
+        self.nodes: dict[str, workspace.WorkspaceNode] = {}
         for node in ws.nodes:
             self.nodes[node.id] = node
             if node.data.title == "Optimizer":
@@ -279,8 +282,8 @@ class ModelBuilder:
         assert len(optimizers) == 1, f"More than one optimizer found: {optimizers}"
         [self.optimizer] = optimizers
         self.dependencies = {n.id: [] for n in ws.nodes}
-        self.in_edges = {}
-        self.out_edges = {}
+        self.in_edges: dict[str, dict[str, list[(str, str)]]] = {}
+        self.out_edges: dict[str, dict[str, list[(str, str)]]] = {}
         repeats = []
         for e in ws.edges:
             if self.nodes[e.target].data.title == "Repeat":
@@ -367,21 +370,12 @@ class ModelBuilder:
         t = node.data.title
         op = self.catalog[t]
         p = op.convert_params(node.data.params)
-        inputs = {}
-        for n in self.in_edges.get(node_id, []):
-            for b, h in self.in_edges[node_id][n]:
-                i = _to_id(b, h)
-                inputs[n] = i
-        outputs = {}
-        for out in self.out_edges.get(node_id, []):
-            i = _to_id(node_id, out)
-            outputs[out] = i
         match t:
             case "Repeat":
                 if node_id.startswith("END "):
                     repeat_id = node_id.removeprefix("END ")
                     start_id = f"START {repeat_id}"
-                    print(f"repeat {repeat_id} ending")
+                    [last_output] = self.in_edges[node_id]["input"]
                     after_start = self.all_downstream(start_id)
                     after_end = self.all_downstream(node_id)
                     before_end = self.all_upstream(node_id)
@@ -390,28 +384,64 @@ class ModelBuilder:
                     assert affected_nodes == repeated_nodes, (
                         f"edges leave repeated section '{repeat_id}':\n{affected_nodes - repeated_nodes}"
                     )
-                    for n in repeated_nodes:
-                        print(f"repeating {n}")
+                    repeated_layers = [e for e in self.layers if e._origin_id in repeated_nodes]
+                    for i in range(p["times"] - 1):
+                        # Copy repeat section's output to repeat section's input.
+                        self.layers.append(
+                            self.empty_layer(
+                                node_id,
+                                inputs=[_to_id(*last_output)],
+                                outputs=[_to_id(start_id, "output")],
+                            )
+                        )
+                        # Repeat the layers in the section.
+                        for layer in repeated_layers:
+                            if p["same_weights"]:
+                                self.layers.append(
+                                    Layer(
+                                        layer.module,
+                                        shapes=layer.shapes,
+                                        _origin_id=layer._origin_id,
+                                        _inputs=layer._inputs,
+                                        _outputs=layer._outputs,
+                                    )
+                                )
+                            else:
+                                self.run_node(layer._origin_id)
+                self.layers.append(self.run_op(node_id, op, p))
             case "Optimizer" | "Input: tensor" | "Input: graph edges" | "Input: sequential":
                 return
-        layer = self.run_op(op, p, inputs, outputs)
-        layer._origin_id = node_id
-        self.layers.append(layer)
+            case _:
+                self.layers.append(self.run_op(node_id, op, p))
 
-    def run_op(self, op, params, inputs: dict[str, str], outputs: dict[str, str]) -> Layer:
+    def run_op(self, node_id: str, op: ops.Op, params) -> Layer:
         """Returns the layer produced by this op."""
-        op_inputs = [
-            TensorRef(inputs[i], shape=self.sizes.get(inputs[i], 1)) for i in op.inputs.keys()
-        ]
+        inputs = [_to_id(*i) for n in op.inputs for i in self.in_edges[node_id][n]]
+        outputs = [_to_id(node_id, n) for n in op.outputs]
+        layer = self.empty_layer(node_id, inputs, outputs)
         if op.func != ops.no_op:
-            layer = op.func(*op_inputs, **params)
-        else:
-            layer = Layer(torch.nn.Identity(), shapes=[i.shape for i in op_inputs])
-        layer._inputs = op_inputs
-        layer._outputs = []
-        for o, shape in zip(op.outputs.keys(), layer.shapes):
-            layer._outputs.append(TensorRef(outputs[o], shape=shape))
-            self.sizes[outputs[o]] = shape
+            op_layer = op.func(*layer._inputs, **params)
+            layer.module = op_layer.module
+            layer.shapes = op_layer.shapes
+            for o in layer._outputs:
+                self.sizes[o._id] = o.shape
+        return layer
+
+    def empty_layer(self, id: str, inputs: list[str], outputs: list[str]) -> Layer:
+        """Creates an identity layer. Assumes that outputs have the same shapes as inputs."""
+        layer_inputs = [TensorRef(i, shape=self.sizes.get(i, 1)) for i in inputs]
+        layer_outputs = []
+        for i, o in zip(inputs, outputs):
+            shape = self.sizes.get(i, 1)
+            layer_outputs.append(TensorRef(o, shape=shape))
+            self.sizes[o] = shape
+        layer = Layer(
+            torch.nn.Identity(),
+            shapes=[self.sizes[o._id] for o in layer_outputs],
+            _inputs=layer_inputs,
+            _outputs=layer_outputs,
+            _origin_id=id,
+        )
         return layer
 
     def build_model(self) -> ModelConfig:
@@ -462,7 +492,6 @@ class ModelBuilder:
         p = op.convert_params(self.nodes[self.optimizer].data.params)
         o = getattr(torch.optim, p["type"].name)
         cfg["optimizer"] = o(cfg["model"].parameters(), lr=p["lr"])
-        print(cfg)
         return ModelConfig(**cfg)