(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 487 | |
| 488 | @register_lower_rule(mops.IndexingMultiAxisVec) |
| 489 | def vec_indexing_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 490 | assert len(ctx.param["items"]) == 1 |
| 491 | axis, _, _, _, is_index = ctx.param["items"][0] |
| 492 | assert is_index |
| 493 | inp, indices = args[0], args[1] |
| 494 | indices = convert_negative_index(indices, inp.shape[axis]) |
| 495 | indices = indices.reshape(indices.shape + (1,)) |
| 496 | slices_size = tuple( |
| 497 | (inp.shape[i] if i != axis else 1 for i in range(len(inp.shape))) |
| 498 | ) |
| 499 | return xla_gather( |
| 500 | inp, |
| 501 | indices, |
| 502 | slices_size, |
| 503 | offset_dims=tuple(i for i in range(len(inp.shape)) if i != axis), |
| 504 | collapsed_slice_dims=(axis,), |
| 505 | start_index_map=(axis,), |
| 506 | ) |
| 507 | |
| 508 | |
| 509 | @register_lower_rule(mops.IndexingIncrMultiAxisVec) |
nothing calls this directly
no test coverage detected