Use scipy.sparse in testing to avoid creating many matrix densifying functions

Currently, we have many functions like csr_densify5x5 and csc_densify8x8 in many of our tests.

Since many of our test files use their own JIT engine fixtures whose scope is just the test file, we have to redundantly have this same densifying code in many test files.

It'd be nicer to have a Python function that takes in a MLIRSparseTensor and returns a np.ndarray.

It might look something like this:

import numpy as np
import scipy.sparse as ss
from mlir_graphblas.sparse_utils import MLIRSparseTensor

def densify_csc(tensor_csc: MLIRSparseTensor) -> np.ndarray:
    ss_tensor_csc = ss.csc_matrix(
        (tensor_csc.values, tensor_csc.indices[1], tensor_csc.pointers[1]),
        shape=(
            tensor_csc.get_dimsize(0),
            tensor_csc.get_dimsize(1)
        )
    )
    return ss_tensor_csc.toarray()

mlir_graphblas/tests/jit_engine_test_utils.py would be an appropriate place to put it.

We'd need one for CSR as well. It might be worth investigating the same for sparse vectors.

This would allow us to get rid of all the densifying functions in our JIT engine tests. This would also allow us to have one densifying function for all dtypes.

This would require us to have scipy in our development environments, which would require us to change some of the environment yml files so that our CI jobs don't fail.

We could use this code to densify vectors:

def densify_vector(sparse_vec: MLIRSparseTensor) -> np.ndarray:
    (vec_length,) = sparse_vec.shape
    dense_vec = np.zeros(vec_length, dtype=sparse_vec.values.dtype)
    dense_vec[sparse_vec.indices] = sparse_vec.values
    return dense_vec

metagraph-dev / mlir-graphblas

Use scipy.sparse in testing to avoid creating many matrix densifying functions #122