( # noqa: C901
cls, ep: ExportedProgram, head: Node
)
| 80 | |
| 81 | @classmethod |
| 82 | def maybe_create( # noqa: C901 |
| 83 | cls, ep: ExportedProgram, head: Node |
| 84 | ) -> Optional["IndexCopyHandler"]: |
| 85 | index_copy_node = head |
| 86 | if not match_target(index_copy_node, torch.ops.aten.index_copy.default): |
| 87 | return None |
| 88 | |
| 89 | # index_copy should write to a mutable input/buffer to be an index update. |
| 90 | if (index_copy_node.name not in ep.graph_signature.buffers_to_mutate) and ( |
| 91 | index_copy_node.name not in ep.graph_signature.user_inputs_to_mutate |
| 92 | ): |
| 93 | return None |
| 94 | |
| 95 | # index_copy(dst, axis, indices, update) |
| 96 | if len(index_copy_node.args) != 4: |
| 97 | return None |
| 98 | dst, axis, indices, update = index_copy_node.args |
| 99 | |
| 100 | # axis must be a literal int |
| 101 | if not isinstance(axis, int): |
| 102 | return None |
| 103 | |
| 104 | return cls( |
| 105 | head=index_copy_node, |
| 106 | body=[], |
| 107 | dst=dst, |
| 108 | update=update, |
| 109 | indices=indices, |
| 110 | axis=axis, |
| 111 | ) |
| 112 | |
| 113 | def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: |
| 114 | assert n == self.head |
nothing calls this directly
no test coverage detected