Handle narrow(input, dim, start, length) -> slice(input, dim, start, start+length). This is needed for KV cache updates with dynamic positions where narrow is preferred over slice syntax for better torch.export compatibility.
(P: MLXProgramBuilder, n: Node)
| 1295 | |
| 1296 | @REGISTRY.register(target=[torch.ops.aten.narrow.default]) |
| 1297 | def _narrow_handler(P: MLXProgramBuilder, n: Node) -> Slot: |
| 1298 | """ |
| 1299 | Handle narrow(input, dim, start, length) -> slice(input, dim, start, start+length). |
| 1300 | |
| 1301 | This is needed for KV cache updates with dynamic positions where narrow |
| 1302 | is preferred over slice syntax for better torch.export compatibility. |
| 1303 | """ |
| 1304 | args = P.args(n) |
| 1305 | require_args(args, 4, 4, "aten.narrow") |
| 1306 | require_kwargs(P.kwargs(n), set(), "aten.narrow") |
| 1307 | x, dim, start, length = args |
| 1308 | out = P.make_or_get_slot(n) |
| 1309 | |
| 1310 | # Convert narrow (start, length) to slice (start, end) |
| 1311 | # The end is start + length |
| 1312 | start_iov = P.to_int_or_vid(start) |
| 1313 | length_iov = P.to_int_or_vid(length) |
| 1314 | |
| 1315 | # For stop = start + length, we need to emit an ADD_SCALAR if either is a Vid |
| 1316 | if isinstance(start_iov, IntOrVid) and start_iov.vid is not None: |
| 1317 | # start is a Vid, need to add at runtime |
| 1318 | if isinstance(length_iov, IntOrVid) and length_iov.vid is not None: |
| 1319 | # Both are Vids - emit add to compute stop |
| 1320 | _, stop_slot = P.make_tmp_value_slot() |
| 1321 | stop_vid = P.slot_to_vid(stop_slot) |
| 1322 | P.emit( |
| 1323 | AddIntNode( |
| 1324 | a=start_iov.vid, |
| 1325 | b=length_iov.vid, |
| 1326 | out=stop_vid, |
| 1327 | ) |
| 1328 | ) |
| 1329 | stop_iov = IntOrVid(int64=None, vid=stop_vid) |
| 1330 | else: |
| 1331 | # start is Vid, length is int - emit add scalar |
| 1332 | _, stop_slot = P.make_tmp_value_slot() |
| 1333 | stop_vid = P.slot_to_vid(stop_slot) |
| 1334 | P.emit( |
| 1335 | AddIntNode( |
| 1336 | a=start_iov.vid, |
| 1337 | b=( |
| 1338 | length_iov.int64 |
| 1339 | if isinstance(length_iov, IntOrVid) |
| 1340 | else length_iov |
| 1341 | ), |
| 1342 | out=stop_vid, |
| 1343 | ) |
| 1344 | ) |
| 1345 | stop_iov = IntOrVid(int64=None, vid=stop_vid) |
| 1346 | elif isinstance(length_iov, IntOrVid) and length_iov.vid is not None: |
| 1347 | # length is Vid, start is int - emit add scalar |
| 1348 | start_val = start_iov.int64 if isinstance(start_iov, IntOrVid) else start_iov |
| 1349 | _, stop_slot = P.make_tmp_value_slot() |
| 1350 | stop_vid = P.slot_to_vid(stop_slot) |
| 1351 | P.emit( |
| 1352 | AddIntNode( |
| 1353 | a=length_iov.vid, |
| 1354 | b=start_val, |
nothing calls this directly
no test coverage detected