Does the BCCO sparse matrix not support the operation power?

[yank666]

'jnp.power' or operator '**' will raise error "sparse rule for integer_pow is not implemented"
jnp.linalg.matrix_power is usable ， but not always get correct result compared with numpy
For example, the following code will produce incorrect results“

def test_spnorm(M):
    reuslt_dot = M @ M.transpose()
    # sqnorm = jnp.power(reuslt_dot, jnp.ones(reuslt_dot.shape()) *2)                     ##error "sparse rule for integer_pow is not implemented"
    # sqnorm = reuslt_dot ** 2                                     ##error "sparse rule for integer_pow is not implemented"
    sqnorm = jnp.linalg.matrix_power(reuslt_dot, 2)  ##the result not always right
    result = sqnorm.sum()   
    return result

np.random.seed(12)
warm_up = 1
rows = 2  #1000
cols = int(1e+1)  #1e+9
density = 0.5
_rows = np.random.randint(0, int(rows), size=(int(rows * cols * density)))
_cols = np.random.randint(0, int(cols), size=(int(rows * cols * density)))
_data = np.random.rand(int(rows * cols * density))
sp_mat = sp.coo_matrix((_data, (_rows, _cols)), shape=(rows, cols), dtype=np.float64)
mat_bcoo = jsparse.BCOO.from_scipy_sparse(sp_mat)
mat_bcoo = mat_bcoo.sort_indices()
func_jit_jsparse = jsparse.sparsify(test_spnorm)
jax_result = func_jit_jsparse(mat_bcoo, mat_bcoo.transpose()).block_until_ready()

[jakevdp]

Currently no, the general power operation is not supported for BCOO matrices. This is because in general, jnp.power(0, x) does not equal zero, and so the operation is not safe on sparse matrices.

That said, your question made me realize that we could implement lax.integer_pow for sparse matrices, because the exponent is static and we can limit it to positive values where integer_pow(0, x) does equal zero.

[jakevdp]

#14855 implements integer power for sparse matrices.

[yank666]

Thanks for your response sincerely.

[yank666]

by the way , is BCCO sparse matrix also not support jnp.max , sub between dense and sparse?

[jakevdp]

No, because in general this leads to a dense result. If that is what you want, you can convert the sparse matrix to dense and then compute the max.

[yank666]

OK，thx