The data_node was produced (or last accessed as via split) in scope earlier than the current one, traverse the scopes between that level and current one, and insert split nodes. Parameters ---------- data_node : DataNode The data node that we want to use
(self, data_node, stack_level)
| 240 | raise ValueError(f"{data_node} was not produced within this trace.") |
| 241 | |
| 242 | def _realize_split(self, data_node, stack_level): |
| 243 | """The data_node was produced (or last accessed as via split) in scope earlier than the |
| 244 | current one, traverse the scopes between that level and current one, and insert split nodes. |
| 245 | |
| 246 | Parameters |
| 247 | ---------- |
| 248 | data_node : DataNode |
| 249 | The data node that we want to use in the current scope. |
| 250 | stack_level : int |
| 251 | Stack level where the data_node was last "seen". |
| 252 | |
| 253 | Returns |
| 254 | ------- |
| 255 | DataNode |
| 256 | New node that can be used in current branch and scope. |
| 257 | """ |
| 258 | assert 0 <= stack_level and stack_level < self.stack_depth() - 1 |
| 259 | produced_data_node = self._stack[stack_level].get(data_node) |
| 260 | bottom = self._stack[: stack_level + 1] |
| 261 | top = self._stack[stack_level + 1 :] |
| 262 | self._stack = bottom |
| 263 | while top: |
| 264 | current_entry = top.pop(0) |
| 265 | predicate = current_entry.predicate |
| 266 | |
| 267 | # Do not automatically register the outputs in the current scope, we track them below |
| 268 | # in their respective branches. |
| 269 | logging.log( |
| 270 | 9, |
| 271 | ( |
| 272 | f"{self._indent()}[IF] Inserting split" |
| 273 | f" at {self.stack_depth() - 1}:" |
| 274 | f" split({produced_data_node}, predicate={predicate}." |
| 275 | ), |
| 276 | ) |
| 277 | self._is_registration_allowed = False |
| 278 | true_node, false_node = fn._conditional.split( |
| 279 | produced_data_node, predicate=predicate, _if_stmt=True |
| 280 | ) |
| 281 | self._is_registration_allowed = True |
| 282 | |
| 283 | # Record the result of splitting the `data_node` that we are trying to look up |
| 284 | # (short-cut for consecutive lookups) |
| 285 | current_entry.add_split(data_node, produced_data_node, true_node, false_node) |
| 286 | if current_entry.branch == _Branch.TrueBranch: |
| 287 | produced_data_node = true_node |
| 288 | else: |
| 289 | produced_data_node = false_node |
| 290 | self._stack.append(current_entry) |
| 291 | return produced_data_node |
| 292 | |
| 293 | def preprocess_input(self, data_node): |
| 294 | """Process the DataNode that is an input to an operator call. Detect if the DataNode was |