import itertools


def _get_indices(A):
    """Gets the index for each element in the array."""
    dim_ranges = [range(size) for size in A.shape]
    if len(dim_ranges) == 1:
        return dim_ranges[0]
    return itertools.product(*dim_ranges)