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)
def _get_index_nd ( self , key ) : def _is_advanced_index ( index ) : 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' need_broadcast = ( len ( key ) != 1 ) advanced_indices = [ ] basic_indices = [ ] advanced_index_bshape = None 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 ) 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 ) 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 = [ ] 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 ) if need_broadcast : broadcasted_indices = [ ] idx_rshape = [ 1 ] * len ( index_bshape ) if advanced_indices_adjacent : advanced_index_bshape_start = advanced_indices [ 0 ] 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 for dim_size in indices [ i ] . shape [ : : - 1 ] : k -= 1 idx_rshape [ k ] = dim_size else : if i < advanced_indices [ 0 ] : idx_rshape [ i ] = indices [ i ] . shape [ 0 ] elif i > advanced_indices [ - 1 ] : 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 ] ) ) broadcasted_indices . append ( indices [ i ] . reshape ( tuple ( idx_rshape ) ) . broadcast_to ( index_bshape ) ) 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 broadcasted_indices . append ( indices [ i ] . reshape ( tuple ( idx_rshape ) ) . broadcast_to ( index_bshape ) ) for j , _ in enumerate ( idx_rshape ) : idx_rshape [ j ] = 1 indices = broadcasted_indices return op . stack ( * indices )
Returns an index array for use in scatter_nd and gather_nd.