Handle aten.topk - top-k elements along an axis. Decomposes into: partition → slice → sort → reverse (for values) argpartition → slice → gather → argsort → reverse → reorder (for indices) torch.topk returns (values, indices) sorted descending.
(P: MLXProgramBuilder, n: Node)
| 3990 | |
| 3991 | @REGISTRY.register(target=[torch.ops.aten.topk.default]) |
| 3992 | def _topk_handler(P: MLXProgramBuilder, n: Node) -> Slot: |
| 3993 | """Handle aten.topk - top-k elements along an axis. |
| 3994 | |
| 3995 | Decomposes into: partition → slice → sort → reverse (for values) |
| 3996 | argpartition → slice → gather → argsort → reverse → reorder (for indices) |
| 3997 | |
| 3998 | torch.topk returns (values, indices) sorted descending. |
| 3999 | """ |
| 4000 | args = P.args(n) |
| 4001 | require_args(args, 2, 5, "aten.topk") |
| 4002 | require_kwargs(P.kwargs(n), set(), "aten.topk") |
| 4003 | x = args[0] |
| 4004 | k = args[1] |
| 4005 | dim = args[2] if len(args) > 2 else -1 |
| 4006 | |
| 4007 | output_slots = P.make_or_get_slots(n) |
| 4008 | values_slot, indices_slot = output_slots |
| 4009 | |
| 4010 | used = used_getitem_indices(n) |
| 4011 | |
| 4012 | # Get dim size from input metadata for forward slice stop |
| 4013 | x_meta = n.args[0].meta.get("val") |
| 4014 | if x_meta is None: |
| 4015 | raise ValueError("Input tensor metadata not found for topk") |
| 4016 | norm_axis = dim if dim >= 0 else dim + len(x_meta.shape) |
| 4017 | dim_size = x_meta.shape[norm_axis] |
| 4018 | |
| 4019 | # Compute -k for partition index and forward slice start |
| 4020 | if isinstance(k, int): |
| 4021 | neg_k = P.to_int_or_vid(-k) |
| 4022 | # Reverse slice: start=k-1, stop=-(k+1) on the k-sized sliced tensor |
| 4023 | rev_start = P.to_int_or_vid(k - 1) |
| 4024 | rev_stop = P.to_int_or_vid(-(k + 1)) |
| 4025 | else: |
| 4026 | # k is dynamic — emit neg_k = k * -1 at runtime |
| 4027 | _, neg_k_slot = P.make_tmp_value_slot() |
| 4028 | P.emit( |
| 4029 | MultiplyIntNode( |
| 4030 | a=P.to_int_or_vid(k), |
| 4031 | b=IntOrVid.from_literal(-1), |
| 4032 | out=P.slot_to_vid(neg_k_slot), |
| 4033 | ) |
| 4034 | ) |
| 4035 | neg_k = P.to_int_or_vid(neg_k_slot) |
| 4036 | # rev_start = k - 1 |
| 4037 | _, rev_start_slot = P.make_tmp_value_slot() |
| 4038 | P.emit( |
| 4039 | AddIntNode( |
| 4040 | a=P.to_int_or_vid(k), |
| 4041 | b=IntOrVid.from_literal(-1), |
| 4042 | out=P.slot_to_vid(rev_start_slot), |
| 4043 | ) |
| 4044 | ) |
| 4045 | rev_start = P.to_int_or_vid(rev_start_slot) |
| 4046 | # rev_stop = -(k + 1) = neg_k - 1 |
| 4047 | _, rev_stop_slot = P.make_tmp_value_slot() |
| 4048 | P.emit( |
| 4049 | AddIntNode( |
nothing calls this directly
no test coverage detected