(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]])
| 364 | |
| 365 | @register_lower_rule("ArgsortBackward") |
| 366 | def argsort_backward_lower(ctx, *args: Union[HLOTensor, Sequence[HLOTensor]]): |
| 367 | assert ( |
| 368 | len(args) == 3 and len(ctx.vars_in) == 3 and len(ctx.vars_out) == 1 |
| 369 | ), f"{len(args)}, {len(ctx.vars_in)}, {len(ctx.vars_out)}" |
| 370 | dy, idx, x = args[0], args[1], args[2] |
| 371 | if _shape_equal(x.shape, dy.shape): |
| 372 | # for argsort backward |
| 373 | _, dx = _sort_according_to_key( |
| 374 | idx, dy, axis=-1, descending=False, is_stable=True |
| 375 | ) |
| 376 | else: |
| 377 | # for topk backward, only support axis=-1 and the dx is 2d tensor |
| 378 | dx = fill(0, ctx.vars_out[0].shape, ctx.vars_out[0].dtype) |
| 379 | expander = iota(np.int32, idx.shape, dimension=0) |
| 380 | idx = expand_dims(idx, -1) |
| 381 | expander = expand_dims(expander, -1) |
| 382 | idx = concat([expander, idx], -1) |
| 383 | |
| 384 | dx = xla_scatter( |
| 385 | dx, |
| 386 | idx, |
| 387 | dy, |
| 388 | update_window_dims=tuple(), |
| 389 | inserted_window_dims=(0, 1), |
| 390 | scattered_dims_to_operand_dims=(0, 1), |
| 391 | unique_indices=True, |
| 392 | ) |
| 393 | return dx |
| 394 | |
| 395 | |
| 396 | def topk(inp, k, descending=True, kth_only=False, no_sort=False): |
nothing calls this directly
no test coverage detected