Split repeat boxes into start/end boxes.

examples/Model definition

@@ -34,6 +34,20 @@
       "sourceHandle": "loss",
       "target": "Optimizer 2",
       "targetHandle": "loss"
+    },
+    {
+      "id": "Activation 2 Repeat 1",
+      "source": "Activation 2",
+      "sourceHandle": "x",
+      "target": "Repeat 1",
+      "targetHandle": "input"
+    },
+    {
+      "id": "Repeat 1 Linear 1",
+      "source": "Repeat 1",
+      "sourceHandle": "output",
+      "target": "Linear 1",
+      "targetHandle": "x"
     }
   ],
   "env": "PyTorch model",
@@ -42,6 +56,7 @@
       "data": {
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {},
           "name": "Input: embedding",
@@ -75,6 +90,7 @@
       "data": {
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
             "x": {
@@ -126,6 +142,7 @@
       "data": {
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
             "x": {
@@ -174,6 +191,7 @@
       "data": {
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {},
           "name": "Input: label",
@@ -209,6 +227,7 @@
         "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
             "x": {
@@ -243,10 +262,6 @@
               }
             }
           },
-          "position": {
-            "x": 419.0,
-            "y": 396.0
-          },
           "type": "basic"
         },
         "params": {
@@ -271,6 +286,7 @@
         "collapsed": null,
         "display": null,
         "error": null,
+        "input_metadata": null,
         "meta": {
           "inputs": {
             "loss": {
@@ -307,10 +323,6 @@
               }
             }
           },
-          "position": {
-            "x": 526.0,
-            "y": 116.0
-          },
           "type": "basic"
         },
         "params": {
@@ -329,6 +341,72 @@
       },
       "type": "basic",
       "width": 200.0
+    },
+    {
+      "data": {
+        "__execution_delay": 0.0,
+        "collapsed": null,
+        "display": null,
+        "error": null,
+        "input_metadata": null,
+        "meta": {
+          "inputs": {
+            "input": {
+              "name": "input",
+              "position": "top",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "name": "Repeat",
+          "outputs": {
+            "output": {
+              "name": "output",
+              "position": "bottom",
+              "type": {
+                "type": "tensor"
+              }
+            }
+          },
+          "params": {
+            "same_weights": {
+              "default": true,
+              "name": "same_weights",
+              "type": {
+                "type": "<class 'bool'>"
+              }
+            },
+            "times": {
+              "default": 1.0,
+              "name": "times",
+              "type": {
+                "type": "<class 'int'>"
+              }
+            }
+          },
+          "position": {
+            "x": 386.0,
+            "y": 456.0
+          },
+          "type": "basic"
+        },
+        "params": {
+          "same_weights": false,
+          "times": "2"
+        },
+        "status": "planned",
+        "title": "Repeat"
+      },
+      "dragHandle": ".bg-primary",
+      "height": 200.0,
+      "id": "Repeat 1",
+      "position": {
+        "x": -180.0,
+        "y": -90.0
+      },
+      "type": "basic",
+      "width": 200.0
     }
   ]
 }

examples/Model use

@@ -579,54 +579,54 @@
               ],
               "data": [
                 [
-                  "[0.49691743 0.61873293 0.90698647 0.94486356]",
-                  "[1.49691749 1.61873293 1.90698647 1.94486356]",
-                  "[1.4993021488189697, 1.6404846906661987, 1.923316240310669, 1.9422152042388916]"
-                ],
-                [
-                  "[0.56922203 0.98222166 0.76851749 0.28615737]",
-                  "[1.56922197 1.9822216  1.76851749 1.28615737]",
-                  "[1.5835213661193848, 1.9884355068206787, 1.7694181203842163, 1.2917503118515015]"
+                  "[0.67418337 0.79634351 0.23229051 0.71345252]",
+                  "[1.67418337 1.79634356 1.23229051 1.71345258]",
+                  "[1.2304607629776, 1.2535771131515503, 1.4764351844787598, 1.524468183517456]"
                 ],
                 [
-                  "[0.90817457 0.89270043 0.38583666 0.66566533]",
-                  "[1.90817451 1.89270043 1.3858366  1.66566539]",
-                  "[1.9053494930267334, 1.9083378314971924, 1.3998609781265259, 1.6636812686920166]"
+                  "[0.6032477  0.83361369 0.18538666 0.19108021]",
+                  "[1.60324764 1.83361363 1.18538666 1.19108021]",
+                  "[1.0745662450790405, 1.3498694896697998, 1.524811029434204, 1.2336101531982422]"
                 ],
                 [
-                  "[0.72795159 0.79317838 0.27832931 0.96576637]",
-                  "[1.72795153 1.79317832 1.27832937 1.96576643]",
-                  "[1.734963297843933, 1.8026459217071533, 1.2926064729690552, 1.9596911668777466]"
+                  "[0.62569475 0.9881897  0.83639616 0.9828859 ]",
+                  "[1.62569475 1.9881897  1.83639622 1.98288584]",
+                  "[1.3193368911743164, 1.691685438156128, 1.4516379833221436, 1.6486209630966187]"
                 ],
                 [
-                  "[0.04508126 0.76880038 0.80721325 0.62542385]",
-                  "[1.04508126 1.76880038 1.80721331 1.62542391]",
-                  "[1.0830243825912476, 1.7584562301635742, 1.8005754947662354, 1.6277496814727783]"
+                  "[0.63235509 0.70352674 0.96188956 0.46240485]",
+                  "[1.63235509 1.70352674 1.96188951 1.46240485]",
+                  "[1.298614740371704, 1.7463937997817993, 1.4211854934692383, 1.2179659605026245]"
                 ],
                 [
-                  "[0.6032477  0.83361369 0.18538666 0.19108021]",
-                  "[1.60324764 1.83361363 1.18538666 1.19108021]",
-                  "[1.6177492141723633, 1.8144152164459229, 1.1718573570251465, 1.1950569152832031]"
+                  "[0.74064726 0.4155122  0.09800029 0.49930882]",
+                  "[1.74064732 1.4155122  1.09800029 1.49930882]",
+                  "[1.2740169763565063, 1.0242725610733032, 1.413628339767456, 1.288135290145874]"
                 ],
                 [
-                  "[0.15064228 0.03198934 0.25754827 0.51484001]",
-                  "[1.15064228 1.03198934 1.25754833 1.51484001]",
-                  "[1.1556042432785034, 0.9955940246582031, 1.2316606044769287, 1.5150485038757324]"
+                  "[0.34084332 0.73018837 0.54168713 0.91440833]",
+                  "[1.34084332 1.73018837 1.54168713 1.91440833]",
+                  "[1.0640922784805298, 1.3728828430175781, 1.1964272260665894, 1.5791122913360596]"
                 ],
                 [
-                  "[0.48959708 0.48549271 0.32688856 0.356677  ]",
-                  "[1.48959708 1.48549271 1.32688856 1.35667706]",
-                  "[1.4930214881896973, 1.467790961265564, 1.3132573366165161, 1.3589863777160645]"
+                  "[0.80654246 0.08253473 0.74478531 0.71257162]",
+                  "[1.8065424  1.08253479 1.74478531 1.71257162]",
+                  "[1.521227240562439, 1.247929334640503, 1.2715831995010376, 1.1901893615722656]"
                 ],
                 [
-                  "[0.08107251 0.2602725  0.18861133 0.44833237]",
-                  "[1.08107257 1.2602725  1.18861127 1.44833231]",
-                  "[1.102121114730835, 1.2180893421173096, 1.160165548324585, 1.4495322704315186]"
+                  "[0.11560339 0.57495481 0.76535827 0.0391947 ]",
+                  "[1.11560345 1.57495475 1.76535821 1.0391947 ]",
+                  "[0.8016152381896973, 1.6247310638427734, 1.1153241395950317, 0.9691400527954102]"
                 ],
                 [
                   "[0.68094063 0.45189077 0.22661722 0.37354094]",
                   "[1.68094063 1.45189071 1.22661722 1.37354088]",
-                  "[1.6725687980651855, 1.4393560886383057, 1.2169336080551147, 1.3746893405914307]"
+                  "[1.224153995513916, 1.1527836322784424, 1.4024348258972168, 1.204542636871338]"
+                ],
+                [
+                  "[0.79121011 0.54161114 0.69369799 0.1520769 ]",
+                  "[1.79121017 1.54161119 1.69369793 1.15207696]",
+                  "[1.3464093208312988, 1.5594894886016846, 1.5191831588745117, 1.0183898210525513]"
                 ]
               ]
             },
@@ -644,10 +644,6 @@
                   "[0.85706753 0.61447072 0.41741937 0.85147089]",
                   "[1.85706758 1.61447072 1.41741943 1.85147095]"
                 ],
-                [
-                  "[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]"
@@ -688,10 +684,18 @@
                   "[0.11693293 0.49860179 0.55020827 0.88832849]",
                   "[1.11693287 1.49860179 1.55020833 1.88832855]"
                 ],
+                [
+                  "[0.48959708 0.48549271 0.32688856 0.356677  ]",
+                  "[1.48959708 1.48549271 1.32688856 1.35667706]"
+                ],
                 [
                   "[0.50272274 0.54912758 0.17663097 0.79070699]",
                   "[1.50272274 1.54912758 1.17663097 1.79070699]"
                 ],
+                [
+                  "[0.04508126 0.76880038 0.80721325 0.62542385]",
+                  "[1.04508126 1.76880038 1.80721331 1.62542391]"
+                ],
                 [
                   "[0.19908059 0.17570406 0.51475513 0.1893943 ]",
                   "[1.19908059 1.175704   1.51475513 1.18939424]"
@@ -709,13 +713,17 @@
                   "[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]"
                 ],
                 [
                   "[0.88776821 0.51636773 0.30333066 0.32230979]",
                   "[1.88776827 1.51636767 1.30333066 1.32230973]"
                 ],
+                [
+                  "[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]"
@@ -728,10 +736,6 @@
                   "[0.31518555 0.49643308 0.11509258 0.95458382]",
                   "[1.31518555 1.49643302 1.11509252 1.95458388]"
                 ],
-                [
-                  "[0.79121011 0.54161114 0.69369799 0.1520769 ]",
-                  "[1.79121017 1.54161119 1.69369793 1.15207696]"
-                ],
                 [
                   "[0.79423058 0.07138705 0.061777   0.18766576]",
                   "[1.79423058 1.07138705 1.061777   1.1876657 ]"
@@ -764,10 +768,6 @@
                   "[0.98033333 0.97656083 0.38939917 0.81491041]",
                   "[1.98033333 1.97656083 1.38939917 1.81491041]"
                 ],
-                [
-                  "[0.74064726 0.4155122  0.09800029 0.49930882]",
-                  "[1.74064732 1.4155122  1.09800029 1.49930882]"
-                ],
                 [
                   "[0.78956431 0.87284744 0.06880784 0.03455889]",
                   "[1.78956437 1.87284744 1.06880784 1.03455889]"
@@ -804,10 +804,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]"
@@ -816,6 +812,10 @@
                   "[0.77427191 0.21829212 0.12769502 0.74303615]",
                   "[1.77427197 1.21829212 1.12769508 1.74303615]"
                 ],
+                [
+                  "[0.08107251 0.2602725  0.18861133 0.44833237]",
+                  "[1.08107257 1.2602725  1.18861127 1.44833231]"
+                ],
                 [
                   "[0.59812403 0.78395379 0.0291847  0.81814629]",
                   "[1.59812403 1.78395379 1.0291847  1.81814623]"
@@ -856,10 +856,6 @@
                   "[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]"
-                ],
                 [
                   "[0.87285906 0.48354989 0.39394957 0.59456545]",
                   "[1.872859   1.48354983 1.39394951 1.59456539]"
@@ -908,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]"
@@ -924,6 +916,10 @@
                   "[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]"
+                ],
                 [
                   "[0.12024075 0.21342516 0.56858408 0.58644271]",
                   "[1.12024069 1.21342516 1.56858408 1.58644271]"
@@ -933,8 +929,8 @@
                   "[1.91730917 1.22574067 1.09591603 1.33056474]"
                 ],
                 [
-                  "[0.63235509 0.70352674 0.96188956 0.46240485]",
-                  "[1.63235509 1.70352674 1.96188951 1.46240485]"
+                  "[0.49691743 0.61873293 0.90698647 0.94486356]",
+                  "[1.49691749 1.61873293 1.90698647 1.94486356]"
                 ],
                 [
                   "[0.37959969 0.42820001 0.10690689 0.96353984]",
@@ -988,6 +984,10 @@
                   "[0.47870928 0.17129105 0.27300501 0.20634609]",
                   "[1.47870922 1.17129111 1.27300501 1.20634604]"
                 ],
+                [
+                  "[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]"
@@ -1000,7 +1000,7 @@
             }
           },
           "other": {
-            "model": "ModelConfig(model=Sequential(\n  (0) - Linear(in_features=4, out_features=4, bias=True): Input__embedding_1_x -> Linear_1_x\n  (1) - <function leaky_relu at 0x719e0ce23a60>: Linear_1_x -> Activation_2_x\n  (2) - Identity(): Activation_2_x -> Activation_2_x\n), model_inputs=['Input__embedding_1_x'], model_outputs=['Activation_2_x'], loss_inputs=['Input__label_1_y', 'Activation_2_x'], loss=Sequential(\n  (0) - <function mse_loss at 0x719e0ce2d580>: Activation_2_x, Input__label_1_y -> MSE_loss_1_loss\n  (1) - Identity(): MSE_loss_1_loss -> 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__embedding_1_x -> START_Repeat_1_output\n  (1) - Linear(in_features=4, out_features=4, bias=True): START_Repeat_1_output -> Linear_1_x\n  (2) - <function leaky_relu at 0x7e6e1cf2c220>: Linear_1_x -> Activation_2_x\n  (3) - Identity(): Activation_2_x -> END_Repeat_1_output\n  (4) - Identity(): END_Repeat_1_output -> END_Repeat_1_output\n), model_inputs=['Input__embedding_1_x'], model_outputs=['END_Repeat_1_output'], loss_inputs=['Input__label_1_y', 'END_Repeat_1_output'], loss=Sequential(\n  (0) - <function mse_loss at 0x7e6e1cf2dd00>: END_Repeat_1_output, Input__label_1_y -> MSE_loss_1_loss\n  (1) - Identity(): MSE_loss_1_loss -> 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": []
         },
@@ -1036,10 +1036,10 @@
                   ],
                   "loss_inputs": [
                     "Input__label_1_y",
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "outputs": [
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "trained": true
                 },
@@ -1211,10 +1211,10 @@
                   ],
                   "loss_inputs": [
                     "Input__label_1_y",
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "outputs": [
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "trained": false
                 },
@@ -1270,7 +1270,7 @@
           "type": "basic"
         },
         "params": {
-          "epochs": "1001",
+          "epochs": "15",
           "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"column\":\"x\",\"df\":\"df_train\"},\"Input__label_1_y\":{\"column\":\"y\",\"df\":\"df_train\"}}}",
           "model_name": "model"
         },
@@ -1304,7 +1304,6 @@
               },
               "df_test": {
                 "columns": [
-                  "predicted",
                   "x",
                   "y"
                 ]
@@ -1324,10 +1323,10 @@
                   ],
                   "loss_inputs": [
                     "Input__label_1_y",
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "outputs": [
-                    "Activation_2_x"
+                    "END_Repeat_1_output"
                   ],
                   "trained": true
                 },
@@ -1383,9 +1382,9 @@
           "type": "basic"
         },
         "params": {
-          "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"column\":\"x\",\"df\":\"df_test\"}}}",
+          "input_mapping": "{\"map\":{\"Input__embedding_1_x\":{\"df\":\"df_test\",\"column\":\"x\"}}}",
           "model_name": "model",
-          "output_mapping": "{\"map\":{\"Activation_2_x\":{\"column\":\"predicted\",\"df\":\"df_test\"}}}"
+          "output_mapping": "{\"map\":{\"END_Repeat_1_output\":{\"df\":\"df_test\",\"column\":\"predicted\"}}}"
         },
         "status": "done",
         "title": "Model inference"

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

@@ -110,7 +110,7 @@ ops.register_passive_op(
     outputs=[ops.Output(name="output", position="bottom", type="tensor")],
     params=[
         ops.Parameter.basic("times", 1, int),
-        ops.Parameter.basic("same_weights", True, bool),
+        ops.Parameter.basic("same_weights", False, bool),
     ],
 )
 
@@ -207,8 +207,10 @@ def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> Mod
     dependencies = {n.id: [] for n in ws.nodes}
     in_edges = {}
     out_edges = {}
-    # TODO: Dissolve repeat boxes here.
+    repeats = []
     for e in ws.edges:
+        if nodes[e.target].data.title == "Repeat":
+            repeats.append(e.target)
         dependencies[e.target].append(e.source)
         in_edges.setdefault(e.target, {}).setdefault(e.targetHandle, []).append(
             (e.source, e.sourceHandle)
@@ -216,34 +218,78 @@ def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> Mod
         out_edges.setdefault(e.source, {}).setdefault(e.sourceHandle, []).append(
             (e.target, e.targetHandle)
         )
+    # Split repeat boxes into start and end, and insert them into the flow.
+    # TODO: Think about recursive repeats.
+    for repeat in repeats:
+        start_id = f"START {repeat}"
+        end_id = f"END {repeat}"
+        # repeat -> first <- real_input
+        # ...becomes...
+        # real_input -> start -> first
+        first, firsth = out_edges[repeat]["output"][0]
+        [(real_input, real_inputh)] = [
+            k for k in in_edges[first][firsth] if k != (repeat, "output")
+        ]
+        dependencies[first].remove(repeat)
+        dependencies[first].append(start_id)
+        dependencies[start_id] = [real_input]
+        out_edges[real_input][real_inputh] = [
+            k if k != (first, firsth) else (start_id, "input")
+            for k in out_edges[real_input][real_inputh]
+        ]
+        in_edges[start_id] = {"input": [(real_input, real_inputh)]}
+        out_edges[start_id] = {"output": [(first, firsth)]}
+        in_edges[first][firsth] = [(start_id, "output")]
+        # repeat <- last -> real_output
+        # ...becomes...
+        # last -> end -> real_output
+        last, lasth = in_edges[repeat]["input"][0]
+        [(real_output, real_outputh)] = [
+            k for k in out_edges[last][lasth] if k != (repeat, "input")
+        ]
+        del dependencies[repeat]
+        dependencies[end_id] = [last]
+        dependencies[real_output].append(end_id)
+        out_edges[last][lasth] = [(end_id, "input")]
+        in_edges[end_id] = {"input": [(last, lasth)]}
+        out_edges[end_id] = {"output": [(real_output, real_outputh)]}
+        in_edges[real_output][real_outputh] = [
+            k if k != (last, lasth) else (end_id, "output")
+            for k in in_edges[real_output][real_outputh]
+        ]
+    # Walk the graph in topological order.
     sizes = {}
     for k, i in inputs.items():
         sizes[k] = i.shape[-1]
     ts = graphlib.TopologicalSorter(dependencies)
     layers = []
     loss_layers = []
-    in_loss = set()
+    regions: dict[str, set[str]] = {node_id: set() for node_id in dependencies}
     cfg = {}
     used_in_model = set()
     made_in_model = set()
     used_in_loss = set()
     made_in_loss = set()
     for node_id in ts.static_order():
-        node = nodes[node_id]
+        if node_id.startswith("START "):
+            node = nodes[node_id.removeprefix("START ")]
+        elif node_id.startswith("END "):
+            node = nodes[node_id.removeprefix("END ")]
+        else:
+            node = nodes[node_id]
         t = node.data.title
         op = catalog[t]
         p = op.convert_params(node.data.params)
         for b in dependencies[node_id]:
-            if b in in_loss:
-                in_loss.add(node_id)
+            regions[node_id] |= regions[b]
         if "loss" in t:
-            in_loss.add(node_id)
+            regions[node_id].add("loss")
         inputs = {}
         for n in in_edges.get(node_id, []):
             for b, h in in_edges[node_id][n]:
                 i = _to_id(b, h)
                 inputs[n] = i
-                if node_id in in_loss:
+                if "loss" in regions[node_id]:
                     used_in_loss.add(i)
                 else:
                     used_in_model.add(i)
@@ -252,13 +298,13 @@ def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> Mod
             i = _to_id(node_id, out)
             outputs[out] = i
             if inputs:  # Nodes with no inputs are input nodes. Their outputs are not "made" by us.
-                if node_id in in_loss:
+                if "loss" in regions[node_id]:
                     made_in_loss.add(i)
                 else:
                     made_in_model.add(i)
         inputs = types.SimpleNamespace(**inputs)
         outputs = types.SimpleNamespace(**outputs)
-        ls = loss_layers if node_id in in_loss else layers
+        ls = loss_layers if "loss" in regions[node_id] else layers
         match t:
             case "Linear":
                 isize = sizes.get(inputs.x, 1)
@@ -276,6 +322,19 @@ def build_model(ws: workspace.Workspace, inputs: dict[str, torch.Tensor]) -> Mod
                         f"{inputs.x}, {inputs.y} -> {outputs.loss}",
                     )
                 )
+            case "Repeat":
+                ls.append((torch.nn.Identity(), f"{inputs.input} -> {outputs.output}"))
+                sizes[outputs.output] = sizes.get(inputs.input, 1)
+                if node_id.startswith("START "):
+                    regions[node_id].add(("repeat", node_id.removeprefix("START ")))
+                else:
+                    repeat_id = node_id.removeprefix("END ")
+                    print(f"repeat {repeat_id} ending")
+                    regions[node_id].remove(("repeat", repeat_id))
+                    for n in nodes:
+                        r = regions.get(n, set())
+                        if ("repeat", repeat_id) in r:
+                            print(f"repeating {n}")
     cfg["model_inputs"] = list(used_in_model - made_in_model)
     cfg["model_outputs"] = list(made_in_model & used_in_loss)
     cfg["loss_inputs"] = list(used_in_loss - made_in_loss)