Function xla_gather

imperative/python/megengine/xla/rules/tensor.py:345–515

This is xla Gather op wrapper. The biggest difference between xla Gather and mge Gather is that xla Gather gather blocks each time from source. We will use some examples to explain how it works. For more information, you can read the hlo doc: https://github.com/openxla/stableh

(
    src: HLOTensor,
    index: HLOTensor,
    slice_sizes: Sequence[int],
    offset_dims: Sequence[int],
    collapsed_slice_dims: Sequence[int],
    start_index_map: Sequence[int],
    indices_are_sorted: bool = False,
)

def xla_gather(
    src: HLOTensor,
    index: HLOTensor,
    slice_sizes: Sequence[int],
    offset_dims: Sequence[int],
    collapsed_slice_dims: Sequence[int],
    start_index_map: Sequence[int],
    indices_are_sorted: bool = False,
):
    """
    This is xla Gather op wrapper. The biggest difference between xla Gather and mge
    Gather is that xla Gather gather blocks each time from source. 
    
    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/main/docs/spec.md.
    
    ## Case1: What does `src`, `index`, `slice_sizes` mean?
    We have `src` Tensor with shape (4, 5) as below:
        src = [[ 0  1  2  3  4]
               [ 5  6  7  8  9]
               [10 11 12 13 14]
               [15 16 17 18 19]]
    We have `index` Tensor with shape (2,) as below:
        index = [1, 2]
    We set the `slice_sizes` in arglist as (2, 3):
    That means we want to gather a block with shape (2, 3) at `src[1, 2]`, it returns:
        out = [[ 7  8  9]
               [12 13 14]]
    The `out.shape` is (2, 3). The last two dimensions of `out` are same as
    `slice_sizes`.
    
    The last dimension of the `index` will be interpreted as coordinate to indicate
    where we start to gather block. The length of the last dimension of `index` is 2
    because the `src` is 2-d tensor. In this case, we start gather at
    `src[index[0], index[1]]`. 

    Here the `slice_sizes` determine the size of the block to gather. Because the `src`
    is 2-d tensor, the length of `slice_sizes` is also 2. Each element of `slice_sizes`
    corresponds to a dimension of `src`. In this case, we gather block of size (2, 3).

    Factually, the gather operation above is equal to:
        out = src[index[0]:index[0]+slice_sizes[0], index[1]:index[1]+slice_sizes[1]].

    ## Case2: How to do batch gather?
    The source Tensor is the same as above with shape (4, 5).
    The index Tensor with shape (3, 2) as below:
        index = [[0 1]
                 [2 1]
                 [1 2]]
    The `slice_sizes` is also set as (2, 3).
    That means we want to gather 3 block with shape (2, 3) from `src[0, 1]`, 
    `src[2, 1]`, `src[1, 2]`it returns:
        out = [[[ 1  2  3]
                [ 6  7  8]]
               [[11 12 13]
                [16 17 18]]
               [[ 7  8  9]
                [12 13 14]]].