@@ -385,14 +385,6 @@ def _scatter(
385
385
array = [None ]
386
386
digits = math .ceil (math .log10 (len (array ) + 1 ))
387
387
388
- # compile IDs of all body nodes and their gather nodes, which we'll need below
389
- body_node_ids = set ()
390
- for body_node in section .body :
391
- body_node_ids .add (body_node .workflow_node_id )
392
- if isinstance (body_node , Tree .WorkflowSection ):
393
- for gather in body_node .gathers .values ():
394
- body_node_ids .add (gather .workflow_node_id )
395
-
396
388
# for each array element, schedule an instance of the body subgraph
397
389
last_scatter_indices = None
398
390
for i , array_i in enumerate (array ):
@@ -478,12 +470,9 @@ def _gather(gather: Tree.Gather, dependencies: Dict[str, Env.Values]) -> Env.Val
478
470
dep_id_values .append (int (dep_id_fields [- 1 ]))
479
471
assert dep_id_values == list (range (len (dep_ids )))
480
472
481
- # determine if we're ultimately (through nested ops) gathering from a value or a call
482
- leaf = gather
483
- while isinstance (leaf , Tree .Gather ):
484
- leaf = leaf .referee
485
-
486
- # figure out names of the values to gather
473
+ # figure out names of the values to gather, either the name if the referenced decl,
474
+ # or each output of the referenced call.
475
+ leaf = gather .final_referee # follow potential linked list of Gathers for nested sections
487
476
if isinstance (leaf , Tree .Decl ):
488
477
names = [leaf .name ]
489
478
elif isinstance (leaf , Tree .Call ):
0 commit comments