def __process_key(
    arr: DNDarray,
    key: Union[Tuple[int, ...], List[int, ...]],
    return_local_indices: Optional[bool] = False,
    op: Optional[str] = None,
) -> Tuple:
    # """
    # Private method to process the key used for indexing a ``DNDarray`` so that it can be applied to the process-local data, i.e. `key` must be "torch-proof".
    # In a processed key:
    # - any ellipses or newaxis have been replaced with the appropriate number of slice objects
    # - ndarrays and DNDarrays have been converted to torch tensors
    # - the dimensionality is the same as the ``DNDarray`` it indexes
    # This function also manipulates `arr` if necessary, inserting and/or transposing dimensions as indicated by `key`. It calculates the output shape, split axis and balanced status of the indexed array.

    # Parameters
    # ----------
    # arr : DNDarray
    #     The ``DNDarray`` to be indexed
    # key : int, Tuple[int, ...], List[int, ...]
    #     The key used for indexing
    # return_local_indices : bool, optional
    #     Whether to return the process-local indices of the key in the split dimension. This is only possible when the indexing key in the split dimension is ordered e.g. `split_key_is_ordered == 1`. Default: False
    # op : str, optional
    #     The indexing operation that the key is being processed for. Get be "get" for `__getitem__` or "set" for `__setitem__`. Default: "get".

    # Returns
    # -------
    # arr : DNDarray
    #     The ``DNDarray`` to be indexed. Its dimensions might have been modified if advanced, dimensional, broadcasted indexing is used.
    # key : Union(Tuple[Any, ...], DNDarray, np.ndarray, torch.Tensor, slice, int, List[int, ...])
    #     The processed key ready for indexing ``arr``. Its dimensions match the (potentially modified) dimensions of ``arr``.
    #     Note: the key indices along the split axis are LOCAL indices, i.e. refer to the process-local data, if ordered indexing is used. Otherwise, they are GLOBAL indices, referring to the global memory-distributed DNDarray. Communication to extract the non-ordered elements of the input ``DNDarray`` is handled by the ``__getitem__`` function.
    # output_shape : Tuple[int, ...]
    #     The shape of the output ``DNDarray``
    # new_split : int
    #     The new split axis
    # split_key_is_ordered : int
    #     Whether the split key is sorted or ordered. Can be 1: ascending, 0: not ordered, -1: descending order.
    # out_is_balanced : bool
    #     Whether the output ``DNDarray`` is balanced
    # root : int
    #     The root process for the ``MPI.Bcast`` call when single-element indexing along the split axis is used
    # backwards_transpose_axes : Tuple[int, ...]
    #     The axes to transpose the input ``DNDarray`` back to its original shape if it has been transposed for advanced indexing
    # """
    output_shape = list(arr.gshape)
    split_bookkeeping = [None] * arr.ndim
    new_split = arr.split
    arr_is_distributed = False
    if arr.split is not None:
        split_bookkeeping[arr.split] = "split"
        if arr.is_distributed():
            counts, displs = arr.counts_displs()
            arr_is_distributed = True

    advanced_indexing = False
    split_key_is_ordered = 1
    key_is_mask_like = False
    out_is_balanced = False
    root = None
    backwards_transpose_axes = tuple(range(arr.ndim))

    if isinstance(key, list):
        try:
            key = torch.tensor(key, device=arr.larray.device)
        except RuntimeError:
            raise IndexError("Invalid indices: expected a list of integers, got {}".format(key))
    if isinstance(key, (DNDarray, torch.Tensor, np.ndarray)):
        if key.dtype in (bool, uint8, torch.bool, torch.uint8, np.bool_, np.uint8):
            # boolean indexing: shape must be consistent with arr.shape
            key_ndim = key.ndim
            if not tuple(key.shape) == arr.shape[:key_ndim]:
                raise IndexError(
                    "Boolean index of shape {} does not match indexed array of shape {}".format(
                        tuple(key.shape), arr.shape
                    )
                )
            # extract non-zero elements
            try:
                # key is torch tensor
                key = key.nonzero(as_tuple=True)
            except TypeError:
                # key is np.ndarray or DNDarray
                key = key.nonzero()
            key_is_mask_like = True
        else:
            # advanced indexing on first dimension: first dim will expand to shape of key
            output_shape = tuple(list(key.shape) + output_shape[1:])
            # adjust split axis accordingly
            if arr_is_distributed:
                if arr.split != 0:
                    # split axis is not affected
                    split_bookkeeping = [None] * key.ndim + split_bookkeeping[1:]
                    new_split = (
                        split_bookkeeping.index("split")
                        if "split" in split_bookkeeping
                        else None
                    )
                    out_is_balanced = arr.balanced
                else:
                    # split axis is affected
                    if key.ndim > 1:
                        try:
                            key_numel = key.numel()
                        except AttributeError:
                            key_numel = key.size
                        if key_numel == arr.shape[0]:
                            new_split = tuple(key.shape).index(arr.shape[0])
                        else:
                            new_split = key.ndim - 1
                        try:
                            key_split = key[new_split].larray
                            sorted, _ = key_split.sort(stable=True)
                        except AttributeError:
                            key_split = key[new_split]
                            sorted = key_split.sort()
                    else:
                        new_split = 0
                        # assess if key is sorted along split axis
                        try:
                            # DNDarray key
                            sorted, _ = torch.sort(key.larray, stable=True)
                            split_key_is_ordered = torch.tensor(
                                (key.larray == sorted).all(),
                                dtype=torch.uint8,
                                device=key.larray.device,
                            )
                            if key.split is not None:
                                out_is_balanced = key.balanced
                                split_key_is_ordered = (
                                    factories.array(
                                        [split_key_is_ordered],
                                        is_split=0,
                                        device=arr.device,
                                        copy=False,
                                    )
                                    .all()
                                    .astype(types.canonical_heat_types.uint8)
                                    .item()
                                )
                            else:
                                split_key_is_ordered = split_key_is_ordered.item()
                            key = key.larray
                        except AttributeError:
                            # torch or ndarray key
                            try:
                                sorted, _ = torch.sort(key, stable=True)
                            except TypeError:
                                # ndarray key
                                sorted = torch.tensor(np.sort(key), device=arr.larray.device)
                            split_key_is_ordered = (key == sorted).all().item()
                            if not split_key_is_ordered:
                                # prepare for distributed non-ordered indexing: distribute torch/numpy key
                                key = factories.array(key, split=0, device=arr.device).larray
                                out_is_balanced = True
                        if split_key_is_ordered:
                            # extract local key
                            cond1 = key >= displs[arr.comm.rank]
                            cond2 = key < displs[arr.comm.rank] + counts[arr.comm.rank]
                            key = key[cond1 & cond2]
                            if return_local_indices:
                                key -= displs[arr.comm.rank]
                            out_is_balanced = False
            else:
                try:
                    out_is_balanced = key.balanced
                    new_split = key.split
                    key = key.larray
                except AttributeError:
                    # torch or numpy key, non-distributed indexed array
                    out_is_balanced = True
                    new_split = None
            return (
                arr,
                key,
                output_shape,
                new_split,
                split_key_is_ordered,
                key_is_mask_like,
                out_is_balanced,
                root,
                backwards_transpose_axes,
            )

    key = list(key) if isinstance(key, Iterable) else [key]

    # check for ellipsis, newaxis. NB: (np.newaxis is None)==True
    add_dims = sum(k is None for k in key)
    ellipsis = sum(isinstance(k, type(...)) for k in key)
    if ellipsis > 1:
        raise ValueError("indexing key can only contain 1 Ellipsis (...)")
    if ellipsis:
        # key contains exactly 1 ellipsis
        # replace with explicit `slice(None)` for affected dimensions
        # output_shape, split_bookkeeping not affected
        expand_key = [slice(None)] * (arr.ndim + add_dims)
        ellipsis_index = key.index(...)
        ellipsis_dims = arr.ndim - (len(key) - ellipsis - add_dims)
        expand_key[:ellipsis_index] = key[:ellipsis_index]
        expand_key[ellipsis_index + ellipsis_dims :] = key[ellipsis_index + 1 :]
        key = expand_key
    while add_dims > 0:
        # expand array dims: output_shape, split_bookkeeping to reflect newaxis
        # replace newaxis with slice(None) in key
        for i, k in reversed(list(enumerate(key))):
            if k is None:
                key[i] = slice(None)
                arr = arr.expand_dims(i - add_dims + 1)
                output_shape = (
                    output_shape[: i - add_dims + 1] + [1] + output_shape[i - add_dims + 1 :]
                )
                split_bookkeeping = (
                    split_bookkeeping[: i - add_dims + 1]
                    + [None]
                    + split_bookkeeping[i - add_dims + 1 :]
                )
                add_dims -= 1

    # recalculate new_split, transpose_axes after dimensions manipulation
    new_split = split_bookkeeping.index("split") if "split" in split_bookkeeping else None
    transpose_axes, backwards_transpose_axes = tuple(range(arr.ndim)), tuple(range(arr.ndim))
    # check for advanced indexing and slices
    advanced_indexing_dims = []
    advanced_indexing_shapes = []
    lose_dims = 0
    for i, k in enumerate(key):
        if np.isscalar(k) or getattr(k, "ndim", 1) == 0:
            # single-element indexing along axis i
            try:
                output_shape[i], split_bookkeeping[i] = None, None
            except IndexError:
                raise IndexError(
                    f"Too many indices for DNDarray: DNDarray is {arr.ndim}-dimensional, but {len(key)} dimensions were indexed"
                )
            lose_dims += 1
            if i == arr.split:
                key[i], root = arr.__process_scalar_key(
                    k, indexed_axis=i, return_local_indices=return_local_indices
                )
            else:
                key[i], _ = arr.__process_scalar_key(
                    k, indexed_axis=i, return_local_indices=False
                )
        elif isinstance(k, Iterable) or isinstance(k, DNDarray):
            advanced_indexing = True
            advanced_indexing_dims.append(i)
            # work with DNDarrays to assess distribution
            # torch tensors will be extracted in the advanced indexing section below
            k = factories.array(k, device=arr.device, comm=arr.comm, copy=None)
            advanced_indexing_shapes.append(k.gshape)
            if arr_is_distributed and i == arr.split:
                if (
                    not k.is_distributed()
                    and k.ndim == 1
                    and (k.larray == torch.sort(k.larray, stable=True)[0]).all()
                ):
                    split_key_is_ordered = 1
                    out_is_balanced = None
                else:
                    split_key_is_ordered = 0
                    # redistribute key along last axis to match split axis of indexed array
                    k = k.resplit(-1)
                    out_is_balanced = True
            key[i] = k

        elif isinstance(k, slice) and k != slice(None):
            start, stop, step = k.start, k.stop, k.step
            if start is None:
                start = 0
            elif start < 0:
                start += arr.gshape[i]
            if stop is None:
                stop = arr.gshape[i]
            elif stop < 0:
                stop += arr.gshape[i]
            if step is None:
                step = 1
            if step < 0 and start > stop:
                # PyTorch doesn't support negative step as of 1.13
                # Lazy solution, potentially large memory footprint
                # TODO: implement ht.fromiter (implemented in ASSET_ht)
                key[i] = torch.tensor(list(range(start, stop, step)), device=arr.larray.device)
                output_shape[i] = len(key[i])
                split_key_is_ordered = -1
                if arr_is_distributed and new_split == i:
                    if op == "set":
                        # setitem: flip key and keep process-local indices
                        key[i] = key[i].flip(0)
                        cond1 = key[i] >= displs[arr.comm.rank]
                        cond2 = key[i] < displs[arr.comm.rank] + counts[arr.comm.rank]
                        key[i] = key[i][cond1 & cond2]
                        if return_local_indices:
                            key[i] -= displs[arr.comm.rank]
                    else:
                        # getitem: distribute key and proceed with non-ordered indexing
                        key[i] = factories.array(
                            key[i], split=0, device=arr.device, copy=False
                        ).larray
                        out_is_balanced = True
            elif step > 0 and start < stop:
                output_shape[i] = int(torch.tensor((stop - start) / step).ceil().item())
                if arr_is_distributed and new_split == i:
                    split_key_is_ordered = 1
                    out_is_balanced = False
                    local_arr_end = displs[arr.comm.rank] + counts[arr.comm.rank]
                    if stop > displs[arr.comm.rank] and start < local_arr_end:
                        index_in_cycle = (displs[arr.comm.rank] - start) % step
                        if start >= displs[arr.comm.rank]:
                            # slice begins on current rank
                            local_start = start - displs[arr.comm.rank]
                        else:
                            local_start = 0 if index_in_cycle == 0 else step - index_in_cycle
                        if stop <= local_arr_end:
                            # slice ends on current rank
                            local_stop = stop - displs[arr.comm.rank]
                        else:
                            local_stop = local_arr_end
                        key[i] = slice(local_start, local_stop, step)
                    else:
                        key[i] = slice(0, 0)
            elif step == 0:
                raise ValueError("Slice step cannot be zero")
            else:
                key[i] = slice(0, 0)
                output_shape[i] = 0

    if advanced_indexing:
        # adv indexing key elements are DNDarrays: extract torch tensors
        # options: 1. key is mask-like (covers boolean mask as well), 2. adv indexing along split axis, 3. everything else
        # 1. define key as mask-like if each element of key is a DNDarray, and all elements of key are of the same shape, and the advanced-indexing dimensions are consecutive
        key_is_mask_like = (
            all(isinstance(k, DNDarray) for k in key)
            and len(set(k.shape for k in key)) == 1
            and torch.tensor(advanced_indexing_dims).diff().eq(1).all()
        )
        # if split axis is affected by advanced indexing, keep track of non-split dimensions for later
        if arr.is_distributed() and arr.split in advanced_indexing_dims:
            non_split_dims = list(advanced_indexing_dims).copy()
            if arr.split is not None:
                non_split_dims.remove(arr.split)
        # 1. key is mask-like
        if key_is_mask_like:
            key = list(key)
            key_splits = [k.split for k in key]
            if arr.split is not None:
                if not key_splits.count(key_splits[arr.split]) == len(key_splits):
                    if (
                        key_splits[arr.split] is not None
                        and key_splits.count(None) == len(key_splits) - 1
                    ):
                        for i in non_split_dims:
                            key[i] = factories.array(
                                key[i],
                                split=key_splits[arr.split],
                                device=arr.device,
                                comm=arr.comm,
                                copy=None,
                            )
                    else:
                        raise IndexError(
                            f"Indexing arrays must be distributed along the same dimension, got splits {key_splits}."
                        )
                else:
                    # all key_splits must be the same, otherwise raise IndexError
                    if not key_splits.count(key_splits[0]) == len(key_splits):
                        raise IndexError(
                            f"Indexing arrays must be distributed along the same dimension, got splits {key_splits}."
                        )
            # all key elements are now DNDarrays of the same shape, same split axis
        # 2. advanced indexing along split axis
        if arr.is_distributed() and arr.split in advanced_indexing_dims:
            if split_key_is_ordered == 1:
                # extract torch tensors, keep process-local indices only
                k = key[arr.split].larray
                cond1 = k >= displs[arr.comm.rank]
                cond2 = k < displs[arr.comm.rank] + counts[arr.comm.rank]
                k = k[cond1 & cond2]
                if return_local_indices:
                    k -= displs[arr.comm.rank]
                key[arr.split] = k
                for i in non_split_dims:
                    if key_is_mask_like:
                        # select the same elements along non-split dimensions
                        key[i] = key[i].larray[cond1 & cond2]
                    else:
                        key[i] = key[i].larray
            elif split_key_is_ordered == 0:
                # extract torch tensors, any other communication + mask-like case are handled in __getitem__ or __setitem__
                for i in advanced_indexing_dims:
                    key[i] = key[i].larray
            # split_key_is_ordered == -1 not treated here as it is slicing, not advanced indexing
        else:
            # advanced indexing does not affect split axis, return torch tensors
            for i in advanced_indexing_dims:
                key[i] = key[i].larray
        # all adv indexing keys are now torch tensors

        # shapes of adv indexing arrays must be broadcastable
        try:
            broadcasted_shape = torch.broadcast_shapes(*advanced_indexing_shapes)
        except RuntimeError:
            raise IndexError(
                "Shape mismatch: indexing arrays could not be broadcast together with shapes: {}".format(
                    advanced_indexing_shapes
                )
            )
        add_dims = len(broadcasted_shape) - len(advanced_indexing_dims)
        if (
            len(advanced_indexing_dims) == 1
            or list(range(advanced_indexing_dims[0], advanced_indexing_dims[-1] + 1))
            == advanced_indexing_dims
        ):
            # dimensions affected by advanced indexing are consecutive:
            output_shape[
                advanced_indexing_dims[0] : advanced_indexing_dims[0]
                + len(advanced_indexing_dims)
            ] = broadcasted_shape
            if key_is_mask_like:
                # advanced indexing dimensions will be collapsed into one dimension
                if (
                    "split" in split_bookkeeping
                    and split_bookkeeping.index("split") in advanced_indexing_dims
                ):
                    split_bookkeeping[
                        advanced_indexing_dims[0] : advanced_indexing_dims[0]
                        + len(advanced_indexing_dims)
                    ] = ["split"]
                else:
                    split_bookkeeping[
                        advanced_indexing_dims[0] : advanced_indexing_dims[0]
                        + len(advanced_indexing_dims)
                    ] = [None]
            else:
                split_bookkeeping = (
                    split_bookkeeping[: advanced_indexing_dims[0]]
                    + [None] * add_dims
                    + split_bookkeeping[advanced_indexing_dims[0] :]
                )
        else:
            # advanced-indexing dimensions are not consecutive:
            # transpose array to make the advanced-indexing dimensions consecutive as the first dimensions
            non_adv_ind_dims = list(
                i for i in range(arr.ndim) if i not in advanced_indexing_dims
            )
            # keep track of transpose axes order, to be able to transpose back later
            transpose_axes = tuple(advanced_indexing_dims + non_adv_ind_dims)
            arr = arr.transpose(transpose_axes)
            backwards_transpose_axes = tuple(
                torch.tensor(transpose_axes, device=arr.larray.device)
                .argsort(stable=True)
                .tolist()
            )
            # output shape and split bookkeeping
            output_shape = list(output_shape[i] for i in transpose_axes)
            output_shape[: len(advanced_indexing_dims)] = broadcasted_shape
            split_bookkeeping = list(split_bookkeeping[i] for i in transpose_axes)
            split_bookkeeping = [None] * add_dims + split_bookkeeping
            # modify key to match the new dimension order
            key = [key[i] for i in advanced_indexing_dims] + [key[i] for i in non_adv_ind_dims]
            # update advanced-indexing dims
            advanced_indexing_dims = list(range(len(advanced_indexing_dims)))

    # expand key to match the number of dimensions of the DNDarray
    if arr.ndim > len(key):
        key += [slice(None)] * (arr.ndim - len(key))

    key = tuple(key)
    for i in range(output_shape.count(None)):
        lost_dim = output_shape.index(None)
        output_shape.remove(None)
        split_bookkeeping = split_bookkeeping[:lost_dim] + split_bookkeeping[lost_dim + 1 :]
    output_shape = tuple(output_shape)
    new_split = split_bookkeeping.index("split") if "split" in split_bookkeeping else None
    return (
        arr,
        key,
        output_shape,
        new_split,
        split_key_is_ordered,
        key_is_mask_like,
        out_is_balanced,
        root,
        backwards_transpose_axes,
    )