This is xla Scatter op wrapper, which is inverse process of Gather. Please read the docstring of Gather op before reading the below contents. We will use some examples to explain how it works. For more information, you can read the hlo doc: https://github.com/openxla/stablehlo/blob
(
dst: HLOTensor,
index: HLOTensor,
src: HLOTensor,
update_window_dims: Sequence[int],
inserted_window_dims: Sequence[int],
scattered_dims_to_operand_dims: Sequence[int],
indices_are_sorted: bool = False,
unique_indices: bool = False,
reduce_mode: str = None,
)
| 516 | |
| 517 | |
| 518 | def xla_scatter( |
| 519 | dst: HLOTensor, |
| 520 | index: HLOTensor, |
| 521 | src: HLOTensor, |
| 522 | update_window_dims: Sequence[int], |
| 523 | inserted_window_dims: Sequence[int], |
| 524 | scattered_dims_to_operand_dims: Sequence[int], |
| 525 | indices_are_sorted: bool = False, |
| 526 | unique_indices: bool = False, |
| 527 | reduce_mode: str = None, |
| 528 | ): |
| 529 | """ |
| 530 | This is xla Scatter op wrapper, which is inverse process of Gather. Please read the |
| 531 | docstring of Gather op before reading the below contents. |
| 532 | |
| 533 | We will use some examples to explain how it works. For more information, you can |
| 534 | read the hlo doc: https://github.com/openxla/stablehlo/blob/main/docs/spec.md. |
| 535 | |
| 536 | ## Case1: What does `src`, `index`, `dst`, `reduce_mode`, `update_window_dims` do? |
| 537 | We have `src` Tensor with shape (2, 3) as below: |
| 538 | src = [[0 1 2] |
| 539 | [3 4 5]] |
| 540 | We have `index` Tensor with shape (2,): |
| 541 | index = [1, 2] |
| 542 | We have `dst` Tensor is zeros(size=(4, 5)). |
| 543 | |
| 544 | The `update_window_dims` is (0, 1) here. `update_window_dims` determine which dim |
| 545 | of `src` are combined into the scatter block. In this case, it is the dim0 and dim1. |
| 546 | Because the `src` is only 2-d tensor, it means the whole src will be seen as 2-d |
| 547 | block to scatter. |
| 548 | |
| 549 | The last dimension of the `index` is also interpreted as coordinate to indicate |
| 550 | where we start to place the scattered block. The length of the last dimension of |
| 551 | `index` is 2 because the `dst` is 2-d tensor in this case. We place the scattered |
| 552 | block at `dst[index[0], index[1]]`. it returns: |
| 553 | out = [[0 0 0 0 0] |
| 554 | [0 0 0 1 2] |
| 555 | [0 0 3 4 5] |
| 556 | [0 0 0 0 0]] |
| 557 | The `reduce_mode` control how to combine `dst` and the scattered block. |
| 558 | If `reduce_mode` is None, dst[index] = scattered[...]. |
| 559 | If `reduce_mode` is "sum", dst[index] += scattered[...]. |
| 560 | If `reduce_mode` is "min", dst[index] = min(dst[index], scattered[...]). |
| 561 | If `reduce_mode` is "max", dst[index] = max(dst[index], scattered[...]). |
| 562 | |
| 563 | Here we always set `reduce_mode` as `sum`. |
| 564 | |
| 565 | Factually, the scatter operation above is equal to: |
| 566 | scatter_block_shape = src.shape[update_window_dims] # (2, 3) |
| 567 | dst[index[0]:index[0]+scatter_block_shape[0], index[1]:index[1]+scatter_block_shape[1]] += src. |
| 568 | |
| 569 | ## Case2: How to do batch scatter with `update_window_dims`? |
| 570 | We have `src` Tensor with shape (2, 2, 3) as below: |
| 571 | src = [[[ 0 1 2] |
| 572 | [ 3 4 5]] |
| 573 | [[ 6 7 8] |
| 574 | [ 9 10 11]]] |
| 575 | We have `index` Tensor with shape (2, 2): |
no test coverage detected