Returns an index array for use in scatter_nd and gather_nd.

def _get_index_nd(self, key):
        """Returns an index array for use in scatter_nd and gather_nd."""
        def _is_advanced_index(index):
            """The definition of advanced index here includes integers as well, while
            integers are considered as basic index type when the key contains only
            slices and integers."""
            return not isinstance(index, py_slice)

        if isinstance(key, (NDArray, np.ndarray, list, integer_types, py_slice)):
            key = (key,)

        assert isinstance(key, tuple),\
            'index=%s must be a NDArray, or np.ndarray, or list, or tuple ' \
            ' type to use advanced indexing, received type=%s' % (str(key), str(type(key)))

        assert len(key) > 0, "Cannot slice with empty indices"
        shape = self.shape
        assert len(shape) >= len(key),\
            "Slicing dimensions exceeds array dimensions, %d vs %d" % (len(key), len(shape))
        indices = []
        dtype = 'int32'  # index data type passed to gather_nd op
        need_broadcast = (len(key) != 1)
        advanced_indices = []  # include list, NDArray, np.ndarray, integer
        basic_indices = []  # include only slices
        advanced_index_bshape = None  # final advanced index shape
        for i, idx_i in enumerate(key):
            is_advanced_index = True
            if isinstance(idx_i, (np.ndarray, list, tuple)):
                idx_i = array(idx_i, ctx=self.context, dtype=dtype)
                advanced_indices.append(i)
            elif isinstance(idx_i, py_slice):
                start, stop, step = _get_index_range(idx_i.start, idx_i.stop, shape[i], idx_i.step)
                idx_i = arange(start, stop, step, ctx=self.context, dtype=dtype)
                basic_indices.append(i)
                is_advanced_index = False
            elif isinstance(idx_i, integer_types):
                start, stop, step = _get_index_range(idx_i, idx_i+1, shape[i], 1)
                idx_i = arange(start, stop, step, ctx=self.context, dtype=dtype)
                advanced_indices.append(i)
            elif isinstance(idx_i, NDArray):
                if dtype != idx_i.dtype:
                    idx_i = idx_i.astype(dtype)
                advanced_indices.append(i)
            else:
                raise IndexError('Indexing NDArray with index=%s of type=%s is not supported'
                                 % (str(key), str(type(key))))
            if is_advanced_index:
                if advanced_index_bshape is None:
                    advanced_index_bshape = idx_i.shape
                elif advanced_index_bshape != idx_i.shape:
                    need_broadcast = True
                    advanced_index_bshape = _get_broadcast_shape(advanced_index_bshape, idx_i.shape)
            indices.append(idx_i)

        # Get final index shape for gather_nd. See the following reference
        # for determining the output array shape.
        # https://docs.scipy.org/doc/numpy-1.13.0/reference/arrays.indexing.html#combining-advanced-and-basic-indexing  # pylint: disable=line-too-long
        if len(advanced_indices) == 0:
            raise ValueError('Advanced index tuple must contain at least one of the following types:'
                             ' list, tuple, NDArray, np.ndarray, integer, received index=%s' % key)
        # determine the output array's shape by checking whether advanced_indices are all adjacent
        # or separated by slices
        advanced_indices_adjacent = True
        for i in range(0, len(advanced_indices)-1):
            if advanced_indices[i] + 1 != advanced_indices[i+1]:
                advanced_indices_adjacent = False
                break

        index_bshape_list = []  # index broadcasted shape
        if advanced_indices_adjacent:
            for i in range(0, advanced_indices[0]):
                index_bshape_list.extend(indices[i].shape)
                if not need_broadcast and indices[i].shape != advanced_index_bshape:
                    need_broadcast = True
            index_bshape_list.extend(advanced_index_bshape)
            for i in range(advanced_indices[-1]+1, len(indices)):
                if not need_broadcast and indices[i].shape != advanced_index_bshape:
                    need_broadcast = True
                index_bshape_list.extend(indices[i].shape)
        else:
            index_bshape_list.extend(advanced_index_bshape)
            for i in basic_indices:
                index_bshape_list.extend(indices[i].shape)
                if not need_broadcast and indices[i].shape != advanced_index_bshape:
                    need_broadcast = True
        index_bshape = tuple(index_bshape_list)

        # Need to broadcast all ndarrays in indices to the final shape.
        # For example, suppose an array has shape=(5, 6, 7, 8) and
        # key=(slice(1, 5), [[1, 2]], slice(2, 5), [1]).
        # Since key[1] and key[3] are two advanced indices here and they are
        # separated by basic indices key[0] and key[2], the output shape
        # is (1, 2, 4, 3), where the first two elements come from the shape
        # that key[1] and key[3] should broadcast to, which is (1, 2), and
        # the last two elements come from the shape of two basic indices.
        # In order to broadcast all basic and advanced indices to the output shape,
        # we need to reshape them based on their axis. For example, to broadcast key[0],
        # with shape=(4,), we first need to reshape it into (1, 1, 4, 1), and then
        # broadcast the reshaped array to (1, 2, 4, 3); to broadcast key[1], we first
        # reshape it into (1, 2, 1, 1), then broadcast the reshaped array to (1, 2, 4, 3).
        if need_broadcast:
            broadcasted_indices = []
            idx_rshape = [1] * len(index_bshape)
            if advanced_indices_adjacent:
                advanced_index_bshape_start = advanced_indices[0]  # start index of advanced_index_bshape in index_shape
                advanced_index_bshape_stop = advanced_index_bshape_start + len(advanced_index_bshape)
                for i, idx in enumerate(key):
                    if _is_advanced_index(idx):
                        k = advanced_index_bshape_stop
                        # find the reshaped shape for indices[i]
                        for dim_size in indices[i].shape[::-1]:
                            k -= 1
                            idx_rshape[k] = dim_size
                    else:
                        if i < advanced_indices[0]:  # slice is on the left side of advanced indices
                            idx_rshape[i] = indices[i].shape[0]
                        elif i > advanced_indices[-1]:  # slice is on the right side of advanced indices
                            idx_rshape[i-len(key)] = indices[i].shape[0]
                        else:
                            raise ValueError('basic index i=%d cannot be between advanced index i=%d and i=%d'
                                             % (i, advanced_indices[0], advanced_indices[-1]))
                    # broadcast current index to the final shape
                    broadcasted_indices.append(indices[i].reshape(tuple(idx_rshape)).broadcast_to(index_bshape))
                    # reset idx_rshape to ones
                    for j, _ in enumerate(idx_rshape):
                        idx_rshape[j] = 1
            else:
                basic_index_offset = len(advanced_index_bshape)
                for i, idx in enumerate(key):
                    if _is_advanced_index(idx):
                        k = len(advanced_index_bshape)
                        for dim_size in indices[i].shape[::-1]:
                            k -= 1
                            idx_rshape[k] = dim_size
                    else:
                        idx_rshape[basic_index_offset] = indices[i].shape[0]
                        basic_index_offset += 1
                    # broadcast current index to the final shape
                    broadcasted_indices.append(indices[i].reshape(tuple(idx_rshape)).broadcast_to(index_bshape))
                    # reset idx_rshape to ones
                    for j, _ in enumerate(idx_rshape):
                        idx_rshape[j] = 1

            indices = broadcasted_indices
        return op.stack(*indices)