tvm.sparse.column_part_hyb

tvm.sparse.column_part_hyb(num_rows, num_cols, indptr_nd, indices_nd, num_col_parts, buckets)

Partition input CSR matrix by columns and collect rows into buckets according to non zero elements per row.

num_rowsint

Number of rows in the CSR matrix.

num_colsint

Number of columns in the CSR matrix.

indptrNDArray

The indptr array of CSR matrix.

indicesNDArray

The indices array of CSR matrix.

num_col_partsint

Number of column partitions.

bucketsList

The bucket sizes array.

Tuple[List[List[NDArray]]]

The pair of (row_indices, col_indices, mask). row_indices is stored as a list of lists with shape (num_col_parts, len(buckets)), where the innermost element is an NDArray. col_indices and mask are stored in the same way.