fuse as_tensor, such as below: d = prim::dtype(b); c = aten::as_tensor(a, d); -> c = aten::type_as(a, b) */
| 461 | -> c = aten::type_as(a, b) |
| 462 | */ |
| 463 | void FuseAsTensor(Graph* graph, Block* block) { |
| 464 | for (auto it = block->nodes().begin(); it != block->nodes().end();) { |
| 465 | auto* node = *it; |
| 466 | it++; |
| 467 | for (Block* sub_block : node->blocks()) { |
| 468 | FuseAsTensor(graph, sub_block); |
| 469 | } |
| 470 | if (node->kind() == prim::dtype) { |
| 471 | for (auto use : node->output(0)->uses()) { |
| 472 | auto as_tensor = use.user; |
| 473 | Node* typeAs = graph->create(aten::type_as, 1); |
| 474 | typeAs->addInput(as_tensor->input(0)); |
| 475 | typeAs->addInput(node->input(0)); |
| 476 | typeAs->output(0)->copyMetadata(as_tensor->output(0)); |
| 477 | as_tensor->replaceAllUsesWith(typeAs); |
| 478 | as_tensor->removeAllInputs(); |
| 479 | as_tensor->destroy(); |
| 480 | } |
| 481 | } |
| 482 | } |
| 483 | } |
| 484 | |
| 485 | /* |
| 486 | fuse uniform, such as below: |