def compute_graph_array(app: ArrayApplication, ga: GraphArray) -> BlockArray: result_ga: GraphArray = RandomTS(seed=random_seed, max_samples_per_step=1, max_reduction_pairs=1, force_final_action=True).solve(ga) result_ga.grid, result_ga.to_blocks() return BlockArray(result_ga.grid, app.system, result_ga.to_blocks())
def argop(self, op_name: str, arr: BlockArray, axis=None): if len(arr.shape) > 1: raise NotImplementedError( "%s currently supports one-dimensional arrays." % op_name ) if axis is None: axis = 0 assert axis == 0 grid = ArrayGrid(shape=(), block_shape=(), dtype=np.int64.__name__) result = BlockArray(grid, self.cm) reduction_result = None, None for grid_entry in arr.grid.get_entry_iterator(): block_slice: slice = arr.grid.get_slice(grid_entry)[0] block: Block = arr.blocks[grid_entry] syskwargs = { "grid_entry": grid_entry, "grid_shape": arr.grid.grid_shape, "options": {"num_returns": 2}, } reduction_result = self.cm.arg_op( op_name, block.oid, block_slice, *reduction_result, syskwargs=syskwargs ) argoptima, _ = reduction_result result.blocks[()].oid = argoptima return result
def norm(self, X: BlockArray, order=2) -> BlockArray: assert len(X.shape) == 1, "Only vector norms are supported." assert order in (1, 2), "Only order 1 and 2 norms supported." if order == 2: return self.sqrt(X.transpose(defer=True) @ X) else: return self.sum(self.abs(X))
def _stack_copy(self, X): assert len(X.shape) == 1 output_shape = (max(X.shape), max(X.shape)) output_block_shape = (X.block_shape[0], X.block_shape[0]) output_arr_grid = ArrayGrid(output_shape, output_block_shape, X.dtype.__name__) output_block_array = BlockArray(output_arr_grid, self.cm) max_block_rows, max_block_cols = ( output_block_array.blocks.shape[0], output_block_array.blocks.shape[1], ) block_row_index = 0 for i in range(max_block_rows): block_row_index = 0 for j in range(max_block_cols): syskwargs = { "grid_entry": (i, j), "grid_shape": output_arr_grid.grid_shape, } block = output_block_array.blocks[(i, j)] rows, cols = block.shape[0], block.shape[1] output_block_array.blocks[(i, j)].oid = self.cm.triu_copy( X.blocks[block_row_index].oid, rows, cols, syskwargs=syskwargs ) block_row_index += 1 return output_block_array
def reduce( self, op_name: str, X: BlockArray, axis=None, keepdims=False, dtype=None ): res = X.reduce_axis(op_name, axis, keepdims=keepdims) if dtype is not None: res = res.astype(dtype) return res
def arange(self, start_in, shape, block_shape, step=1, dtype=None) -> BlockArray: assert step == 1 if dtype is None: dtype = np.__getattribute__( str(np.result_type(start_in, shape[0] + start_in))) # Generate ranges per block. grid = ArrayGrid(shape, block_shape, dtype.__name__) rarr = BlockArray(grid, self.cm) for _, grid_entry in enumerate(grid.get_entry_iterator()): syskwargs = { "grid_entry": grid_entry, "grid_shape": grid.grid_shape } start = start_in + block_shape[0] * grid_entry[0] entry_shape = grid.get_block_shape(grid_entry) stop = start + entry_shape[0] rarr.blocks[grid_entry].oid = self.cm.arange(start, stop, step, dtype, syskwargs=syskwargs) return rarr
def quantile( self, arr: BlockArray, q: float, interpolation="linear", method="tdigest" ) -> BlockArray: """Compute the q-th quantile of the array elements. Args: arr: BlockArray. q: quantile to compute, which must be between 0 and 1 inclusive. interpolation: interpolation method to use when the desired quantile lies between two data points i < j. also see https://numpy.org/doc/1.20/reference/generated/numpy.quantile.html. also see https://docs.dask.org/en/latest/_modules/dask/array/percentile.html. Returns: Returns the q-th quantile of the array elements. """ # pylint: disable = import-outside-toplevel, unused-import try: import crick except Exception as e: raise Exception( "Unable to import crick. \ Install crick with command 'pip install cython; pip install crick'" ) from e if arr.ndim != 1: raise NotImplementedError("Only 1D 'arr' is currently supported.") if q < 0.0 or q > 1.0: raise ValueError("Quantiles must be in the range [0, 1]") assert interpolation == "linear" assert method == "tdigest" arr_oids = arr.flattened_oids() num_arrs = len(arr_oids) q = [q] t_oids = [] for i, arr_oid in enumerate(arr_oids): syskwargs = { "grid_entry": (i,), "grid_shape": (num_arrs,), "options": {"num_returns": 1}, } t_oids.append(self.cm.tdigest_chunk(arr_oid, syskwargs=syskwargs)) p_oid = self.cm.percentiles_from_tdigest(q, *t_oids, syskwargs=syskwargs) return BlockArray.from_oid(p_oid, (1,), np.float64, self.cm)
def map_bop(self, op_name: str, arr_1: BlockArray, arr_2: BlockArray, out: BlockArray = None, where=True, args=None, kwargs=None) -> BlockArray: # TODO (hme): Move this into BlockArray, and invoke on operator implementations. """ A map, for binary operators, that applies element-wise to every entry of the input arrays. :param op_name: An element-wise binary operator. :param arr_1: A BlockArray. :param arr_2: A BlockArray. :param out: A BlockArray to which the result is written. :param where: An indicator specifying the indices to which op is applied. :param args: Args provided to op. :param kwargs: Keyword args provided to op. :return: A BlockArray. """ if where is not True: raise NotImplementedError("'where' argument is not yet supported.") if args is not None: raise NotImplementedError("'args' is not yet supported.") if not (kwargs is None or len(kwargs) == 0): raise NotImplementedError("'kwargs' is not yet supported.") try: ufunc = np.__getattribute__(op_name) if (op_name.endswith("max") or op_name == "maximum" or op_name.endswith("min") or op_name == "minimum" or op_name.startswith("logical")): rarr = self._broadcast_bop(op_name, arr_1, arr_2) else: result_blocks: np.ndarray = ufunc(arr_1.blocks, arr_2.blocks) rarr = BlockArray.from_blocks(result_blocks, result_shape=None, system=self.system) except Exception as _: rarr = self._broadcast_bop(op_name, arr_1, arr_2) if out is not None: assert out.grid.grid_shape == rarr.grid.grid_shape assert out.shape == rarr.shape assert out.block_shape == rarr.block_shape out.blocks[:] = rarr.blocks[:] rarr = out return rarr
def from_modin(df): # pylint: disable = import-outside-toplevel, protected-access, unidiomatic-typecheck import numpy as np try: from modin.pandas.dataframe import DataFrame from modin.engines.ray.pandas_on_ray.frame.data import PandasOnRayFrame from modin.engines.ray.pandas_on_ray.frame.partition import ( PandasOnRayFramePartition, ) except Exception as e: raise Exception( "Unable to import modin. Install modin with command 'pip install modin'" ) from e assert isinstance( df, DataFrame), "Unexpected dataframe type %s" % str(type(df)) assert isinstance(df._query_compiler._modin_frame, PandasOnRayFrame), "Unexpected dataframe type %s" % str( type(df._query_compiler._modin_frame)) frame: PandasOnRayFrame = df._query_compiler._modin_frame app: ArrayApplication = _instance() system = app.cm # Make sure the partitions are numeric. dtype = frame.dtypes[0] if not array_utils.is_supported(dtype, type_test=True): raise TypeError("%s is not supported." % str(dtype)) for dt in frame.dtypes: if dt != dtype: raise TypeError("Mixed types are not supported (%s != %s).") dtype = np.__getattribute__(str(dtype)) # Convert from Pandas to NumPy. pd_parts = frame._partition_mgr_cls.map_partitions(frame._partitions, lambda df: np.array(df)) grid_shape = len(frame._row_lengths), len(frame._column_widths) shape = (np.sum(frame._row_lengths), np.sum(frame._column_widths)) block_shape = app.get_block_shape(shape, dtype) rows = [] for i in range(grid_shape[0]): cols = [] for j in range(grid_shape[1]): curr_block_shape = (frame._row_lengths[i], frame._column_widths[j]) part: PandasOnRayFramePartition = pd_parts[(i, j)] part.drain_call_queue() ba: BlockArray = BlockArray.from_oid(part.oid, curr_block_shape, dtype, system) cols.append(ba) if grid_shape[1] == 1: row_ba: BlockArray = cols[0] else: row_ba: BlockArray = app.concatenate( cols, axis=1, axis_block_size=block_shape[1]) rows.append(row_ba) result = app.concatenate(rows, axis=0, axis_block_size=block_shape[0]) return result
def permutation(self, size, block_size): shape = (size, ) block_shape = (block_size, ) grid: ArrayGrid = ArrayGrid(shape=shape, block_shape=shape, dtype=np.int64.__name__) ba = BlockArray(grid, self._cm) for grid_entry in ba.grid.get_entry_iterator(): rng_params = list(self._rng.new_block_rng_params()) block: Block = ba.blocks[grid_entry] block.oid = self._cm.permutation(rng_params, size, syskwargs={ "grid_entry": grid_entry, "grid_shape": grid.grid_shape }) return ba.reshape(block_shape=block_shape)
def test_split(app_inst: ArrayApplication): # TODO (hme): Implement a split leveraging block_shape param in reshape op. x = app_inst.array(np.array([1.0, 2.0, 3.0, 4.0]), block_shape=(4,)) syskwargs = { "grid_entry": x.blocks[0].grid_entry, "grid_shape": x.blocks[0].grid_shape, "options": {"num_returns": 2} } res1, res2 = x.cm.split(x.blocks[0].oid, 2, axis=0, transposed=False, syskwargs=syskwargs) ba = BlockArray(ArrayGrid((4,), (2,), x.dtype.__name__), x.cm) ba.blocks[0].oid = res1 ba.blocks[1].oid = res2 assert np.allclose([1.0, 2.0, 3.0, 4.0], ba.get())
def mean(self, X: BlockArray, axis=None, keepdims=False, dtype=None): if X.dtype not in (float, np.float32, np.float64): X = X.astype(np.float64) num_summed = np.product(X.shape) if axis is None else X.shape[axis] res = self.sum(X, axis=axis, keepdims=keepdims) / num_summed if dtype is not None: res = res.astype(dtype) return res
def map_uop( self, op_name: str, arr: BlockArray, out: BlockArray = None, where=True, args=None, kwargs=None, ) -> BlockArray: """A map, for unary operators, that applies to every entry of an array. Args: op_name: An element-wise unary operator. arr: A BlockArray. out: A BlockArray to which the result is written. where: An indicator specifying the indices to which op is applied. args: Args provided to op. kwargs: Keyword args provided to op. Returns: A BlockArray. """ if where is not True: raise NotImplementedError("'where' argument is not yet supported.") args = () if args is None else args kwargs = {} if kwargs is None else kwargs shape = arr.shape block_shape = arr.block_shape dtype = array_utils.get_uop_output_type(op_name, arr.dtype) assert len(shape) == len(block_shape) if out is None: grid = ArrayGrid(shape, block_shape, dtype.__name__) rarr = BlockArray(grid, self.cm) else: rarr = out grid = rarr.grid assert rarr.shape == arr.shape and rarr.block_shape == arr.block_shape for grid_entry in grid.get_entry_iterator(): # TODO(hme): Faster to create ndarray first, # and instantiate block array on return # to avoid instantiating blocks on BlockArray initialization. rarr.blocks[grid_entry] = arr.blocks[grid_entry].uop_map( op_name, args=args, kwargs=kwargs ) return rarr
def triu(self, X: BlockArray): if len(X.shape) == 1: return self.triu(self._stack_copy(X)) elif len(X.shape) == 2: if X.shape[0] == 1: return X diag_meta = array_utils.find_diag_output_blocks(X.blocks, min(X.shape)) output_arr_grid = ArrayGrid(X.shape, X.block_shape, X.dtype.__name__) output_block_array = BlockArray(output_arr_grid, self.cm) visited = dict() total_row_blocks, total_col_blocks = X.blocks.shape[0], X.blocks.shape[1] for block_indices, offset, total_elements in diag_meta: syskwargs = { "grid_entry": block_indices, "grid_shape": output_arr_grid.grid_shape, } output_block_array.blocks[block_indices].oid = self.cm.triu( X.blocks[block_indices].oid, offset, False, total_elements, syskwargs=syskwargs, ) visited[block_indices] = 1 for block_indices, offset, total_elements in diag_meta: row_c, col_c = block_indices[0] + 1, block_indices[1] while row_c < total_row_blocks: syskwargs = { "grid_entry": (row_c, col_c), "grid_shape": output_arr_grid.grid_shape, } if (row_c, col_c) in visited: output_block_array.blocks[(row_c, col_c)].oid = self.cm.triu( output_block_array.blocks[(row_c, col_c)].oid, offset, True, total_elements, syskwargs=syskwargs, ) else: output_block_array.blocks[(row_c, col_c)].oid = self.cm.triu( X.blocks[(row_c, col_c)].oid, offset, True, total_elements, syskwargs=syskwargs, ) visited[(row_c, col_c)] = 1 row_c += 1 for i in range(total_row_blocks): for j in range(total_col_blocks): if (i, j) not in visited: output_block_array.blocks[(i, j)].oid = X.blocks[(i, j)].oid return output_block_array else: raise NotImplementedError()
def inv(app: ArrayApplication, X: BlockArray): # TODO (hme): Implement scalable version. block_shape = X.block_shape assert len(X.shape) == 2 assert X.shape[0] == X.shape[1] single_block = X.shape[0] == X.block_shape[0] and X.shape[1] == X.block_shape[1] if single_block: result = X.copy() else: result = X.reshape(block_shape=X.shape) result.blocks[0, 0].oid = app.cm.inv(result.blocks[0, 0].oid, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1) }) if not single_block: result = result.reshape(block_shape=block_shape) return result
def inv(a: BlockArray): if not a.is_single_block(): warnings.warn( "nums.numpy.linalg.inv is not a scalable implementation. " + ("Input array is %s bytes. " % a.nbytes) + "Abort this operation if input array is too large to " + "execute on a single node." ) return linalg.inv(_instance(), a)
def gradient( self, X: BlockArray, y: BlockArray, mu: BlockArray = None, beta: BlockArray = None, ): if mu is None: mu = self.forward(X) return X.transpose(defer=True) @ (mu - y)
def hessian(self, X: BlockArray, y: BlockArray, mu: BlockArray = None): if mu is None: mu = self.forward(X) dim, block_dim = mu.shape[0], mu.block_shape[0] s = (mu * (self._app.one - mu)).reshape((dim, 1), block_shape=(block_dim, 1)) r = X.transpose(defer=True) @ (s * X) if self._penalty is not None: r += self.hessian_penalty() return r
def column_stack(self, tup, axis_block_size=None): # Based on numpy source. arrays = [] for obj in tup: arr = BlockArray.to_block_array(obj, self.cm) if arr.ndim < 2: arrays.append(self.atleast_2d(arr).T) else: arrays.append(self.atleast_2d(arr)) return self.concatenate(arrays, 1, axis_block_size=axis_block_size)
def train_test_split(*arrays, test_size: Union[int, float] = None, train_size: Union[int, float] = None, random_state: Optional[Union[NumsRandomState, int]] = None, shuffle: bool = True, stratify=None): # pylint: disable = protected-access updated_arrays = [] for array in arrays: updated_arrays.append(_check_array(array)) syskwargs = { "options": { "num_returns": 2 * len(updated_arrays) }, "grid_entry": (0, ), "grid_shape": (1, ), } if random_state is None: rng_params = None else: if isinstance(random_state, int): # It's a seed. random_state: NumsRandomState = instance().random_state( random_state) rng_params = random_state._rng.new_block_rng_params() array_oids = [array.flattened_oids()[0] for array in updated_arrays] result_oids = instance().cm.call("train_test_split", *array_oids, rng_params=rng_params, test_size=test_size, train_size=train_size, shuffle=shuffle, stratify=stratify, syskwargs=syskwargs) # Optimize by computing this directly. shape_dtype_oids = [ instance().cm.shape_dtype(r_oid, syskwargs={ "grid_entry": (0, ), "grid_shape": (1, ) }) for r_oid in result_oids ] shape_dtypes = instance().cm.get(shape_dtype_oids) results = [] for i, r_oid in enumerate(result_oids): shape, dtype = shape_dtypes[i] results.append( BlockArray.from_oid(r_oid, shape=shape, dtype=dtype, cm=instance().cm)) return results
def where(self, condition: BlockArray, x=None, y=None): result_oids = [] shape_oids = [] num_axes = max(1, len(condition.shape)) # Stronger constraint than necessary, but no reason for anything stronger. if x is not None or y is not None: assert x is not None and y is not None assert condition.shape == x.shape == y.shape assert condition.block_shape == x.block_shape == y.block_shape for grid_entry in condition.grid.get_entry_iterator(): block: Block = condition.blocks[grid_entry] block_slice_tuples = condition.grid.get_slice_tuples(grid_entry) roids = self.system.where(block.oid, x, y, block_slice_tuples, syskwargs={ "grid_entry": grid_entry, "grid_shape": condition.grid.grid_shape, "options": { "num_returns": num_axes + 1 } }) block_oids, shape_oid = roids[:-1], roids[-1] shape_oids.append(shape_oid) result_oids.append(block_oids) shapes = self.system.get(shape_oids) result_shape = (np.sum(shapes), ) if result_shape == (0, ): return (self.array(np.array([], dtype=np.int64), block_shape=(0, )), ) # Remove empty shapes. result_shape_pair = [] for i, shape in enumerate(shapes): if np.sum(shape) > 0: result_shape_pair.append((result_oids[i], shape)) result_block_shape = self.compute_block_shape(result_shape, np.int64) result_arrays = [] for axis in range(num_axes): block_arrays = [] for i in range(len(result_oids)): if shapes[i] == (0, ): continue block_arrays.append( BlockArray.from_oid(result_oids[i][axis], shapes[i], np.int64, self.system)) if len(block_arrays) == 1: axis_result = block_arrays[0] else: axis_result = self.concatenate(block_arrays, 0, result_block_shape[0]) result_arrays.append(axis_result) return tuple(result_arrays)
def score(self, X: BlockArray, y: BlockArray, sample_weight: BlockArray = None): _check_array(X, True) _check_array(y, True) if sample_weight is not None: _check_array(sample_weight, True) sample_weight = sample_weight.flattened_oids()[0] r_oid = instance().cm.call_actor_method( self.actor, "score", X.flattened_oids()[0], y.flattened_oids()[0], sample_weight, ) return BlockArray.from_oid(r_oid, shape=(), dtype=float, cm=instance().cm)
def _inv(self, remote_func, kwargs, X: BlockArray): # TODO (hme): Implement scalable version. block_shape = X.block_shape assert len(X.shape) == 2 assert X.shape[0] == X.shape[1] single_block = X.shape[0] == X.block_shape[0] and X.shape[ 1] == X.block_shape[1] if single_block: result = X.copy() else: result = X.reshape(block_shape=X.shape) result.blocks[0, 0].oid = remote_func(result.blocks[0, 0].oid, **kwargs, syskwargs={ "grid_entry": (0, 0), "grid_shape": (1, 1) }) if not single_block: result = result.reshape(block_shape=block_shape) return result
def diag(self, X: BlockArray) -> BlockArray: if len(X.shape) == 1: shape = X.shape[0], X.shape[0] block_shape = X.block_shape[0], X.block_shape[0] grid = ArrayGrid(shape, block_shape, X.dtype.__name__) grid_meta = grid.to_meta() rarr = BlockArray(grid, self.cm) for grid_entry in grid.get_entry_iterator(): syskwargs = { "grid_entry": grid_entry, "grid_shape": grid.grid_shape } if np.all(np.diff(grid_entry) == 0): # This is a diagonal block. rarr.blocks[grid_entry].oid = self.cm.diag( X.blocks[grid_entry[0]].oid, syskwargs=syskwargs) else: rarr.blocks[grid_entry].oid = self.cm.new_block( "zeros", grid_entry, grid_meta, syskwargs=syskwargs) elif len(X.shape) == 2: assert X.shape[0] == X.shape[1], "X must be a square array." assert X.block_shape[0] == X.block_shape[ 1], "block_shape must be square." shape = X.shape[0], block_shape = X.block_shape[0], grid = ArrayGrid(shape, block_shape, X.dtype.__name__) rarr = BlockArray(grid, self.cm) for grid_entry in X.grid.get_entry_iterator(): out_grid_entry = grid_entry[:1] out_grid_shape = grid.grid_shape[:1] syskwargs = { "grid_entry": out_grid_entry, "grid_shape": out_grid_shape } if np.all(np.diff(grid_entry) == 0): # This is a diagonal block. rarr.blocks[out_grid_entry].oid = self.cm.diag( X.blocks[grid_entry].oid, syskwargs=syskwargs) else: raise ValueError("X must have 1 or 2 axes.") return rarr
def array(self, array: Union[np.ndarray, List[float]], block_shape: tuple = None): if not isinstance(array, np.ndarray): if array_utils.is_array_like(array): array = np.array(array) else: raise ValueError( "Unable to instantiate array from type %s" % type(array) ) assert len(array.shape) == len(block_shape) return BlockArray.from_np( array, block_shape=block_shape, copy=False, cm=self.cm )
def read_fs(self, filename: str): meta = self._filesystem.read_meta_fs(filename) addresses = meta["addresses"] grid_meta = meta["grid_meta"] grid = ArrayGrid.from_meta(grid_meta) ba: BlockArray = BlockArray(grid, self._system) for grid_entry in addresses: node_address = addresses[grid_entry] options = {"resources": {node_address: 1.0 / 10**4}} ba.blocks[grid_entry].oid = self._filesystem.read_block_fs( filename, grid_entry, grid_meta, options=options) return ba
def median(self, arr: BlockArray) -> BlockArray: """Compute the median of a BlockArray. Args: a: A BlockArray. Returns: The median value. """ if arr.ndim != 1: raise NotImplementedError("Only 1D 'arr' is currently supported.") a_oids = arr.flattened_oids() if arr.size % 2 == 1: m_oid = self.quickselect(a_oids, arr.size // 2) return BlockArray.from_oid(m_oid, (1, ), arr.dtype, self.cm) else: m0_oid = self.quickselect(a_oids, arr.size // 2 - 1) m0 = BlockArray.from_oid(m0_oid, (1, ), arr.dtype, self.cm) m1_oid = self.quickselect(a_oids, arr.size // 2) m1 = BlockArray.from_oid(m1_oid, (1, ), arr.dtype, self.cm) return (m0 + m1) / 2
def nanvar(self, a: BlockArray, axis=None, ddof=0, keepdims=False, dtype=None): mean = self.nanmean(a, axis=axis, keepdims=True) ss = self.reduce( "nansum", (a - mean) ** self.two, axis=axis, dtype=dtype, keepdims=keepdims ) num_summed = ( self.sum(~a.ufunc("isnan"), axis=axis, dtype=a.dtype, keepdims=keepdims) - ddof ) res = ss / num_summed if dtype is not None: res = res.astype(dtype) return res
def atleast_1d(self, *arys): res = [] for ary in arys: ary = BlockArray.to_block_array(ary, self.cm) if ary.ndim == 0: result = ary.reshape(1) else: result = ary res.append(result) if len(res) == 1: return res[0] else: return res
def top_k(self, arr: BlockArray, k: int, largest=True) -> Tuple[BlockArray, BlockArray]: """Find the `k` largest or smallest elements of a BlockArray. If there are multiple kth elements that are equal in value, then no guarantees are made as to which ones are included in the top k. Args: arr: A BlockArray. k: Number of top elements to return. largest: Whether to return largest or smallest elements. Returns: A tuple containing two BlockArrays, (`values`, `indices`). values: Values of the top k elements, unsorted. indices: Indices of the top k elements, ordered by their corresponding values. """ if arr.ndim != 1: raise NotImplementedError("Only 1D 'arr' is currently supported.") if k <= 0 or arr.size < k: raise IndexError( "'k' must be at least 1 and at most the size of 'arr'.") arr_oids = arr.flattened_oids() if largest: k_oid = self.quickselect(arr_oids, k - 1) k_val = BlockArray.from_oid(k_oid, (1, ), arr.dtype, self.cm) ie_indices = self.where(arr > k_val[0])[0] else: k_oid = self.quickselect(arr_oids, -k) k_val = BlockArray.from_oid(k_oid, (1, ), arr.dtype, self.cm) ie_indices = self.where(arr < k_val[0])[0] eq_indices = self.where(arr == k_val[0])[0] eq_indices_pad = eq_indices[:k - ie_indices.size] axis_block_size = self.compute_block_shape((k, ), int)[0] indices = self.concatenate([ie_indices, eq_indices_pad], 0, axis_block_size) return arr[indices], indices