Method __call__

backends/mlx/patterns.py:113–128

(self, P: MLXProgramBuilder, n: Node)

        )

    def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot:
        assert n == self.head
        dst, update, indices = P.slot_map([self.dst, self.update, self.indices])

        P.emit(
            IndexCopyNode(
                dst=P.slot_to_tid(dst),
                update=P.slot_to_tid(update),
                indices=P.slot_to_tid(indices),
                out=P.slot_to_tid(dst),
                axis=self.axis,
            )
        )

        P.set_slot(n, dst)
        return dst


@REGISTRY.register_pattern(name="ET_KV_CACHE_UPDATE")