from typing import Tuple
import torch


def _get_indexer(shape: Tuple[int]):
    """Return a tuple of Tensors which can be used to index a Tensor.

    The purpose of this function can be better illustrated by an example.
    Suppose ``shape`` is ``[n0, n1, n2]``. Then shape of the three returned
    Tensors will be: [n0, 1, 1], [n1, 1], [n2]. And each of them has elements
    ranging from 0 to n0-1, n1-1 and n2-1 respectively. The returned tuple ``B``
    can be combined with another int64 Tensor ``I`` of shape ``[n0, n1, n2]``
    to access the element of a Tensor ``X`` with shape``[n0, n1, n2, n]`` as
    ``Y=X[B + (I,)]`` so that ``Y[i,j,k] = X[i, j, k, I[i,j,k]]``

    Args:
        shape: The shape of the tensor to be accessed exclusing the last dimension.
    Returns:
        the tuple of index for accessing the tensor.
    """
    ndim = len(shape)
    ones = [1] * ndim
    B = tuple(
        torch.arange(d).reshape(d, *ones[i + 1:]) for i, d in enumerate(shape))
    return B