how to speed jax.experimental.sparse in sparse mat dot with dense vector ,i find it much slower than scipy.sparse and torch.sparse

[yank666]

@jit
def mul_jnp(M, v):
  return 2 * (M @ v )

jnp_sp_mat = sparse.coo_fromdense(dense_mat)
func_sp = sparse.sparsify(mul_jnp)
vec = jnp.asarray(vec, jnp.float32)

vec = jax.device_put(vec)
start = time.time()
for i in range(warm_up):
    result = func_sp(jnp_sp_mat, vec)
    result = sparse.coo_matvec(jnp_sp_mat, vec)
print("jnp sparse * vec  cost :%.4f" % (time.time() - start))    
print(result[0], result.shape)

[jakevdp]

You'll find it faster if you jit compile the sparsified code:

func_sp = sparse.sparsify(mul_jnp)

Also, please see https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code for tips on writing effective microbenchmarks for JAX code, and notes about when and where you can expect JAX to be faster or slower than similar code written in numpy/scipy.

Also, I'd suggest workng with sparse.BCOO rather than sparse.COO; COO only has a handful of operations implemented, and is not compatible with most JAX transformations, including sparsify.

Out of curiosity, I tried a more careful benchmark of JAX vs scipy versions of this operation:

import jax
import jax.numpy as jnp
from jax.experimental import sparse as jsparse
import numpy as np
from scipy import sparse


mat = np.random.rand(1000, 1000)
mat[mat < 0.5] = 0
mat_coo = sparse.coo_matrix(mat)
vec = np.random.rand(1000)

mat_bcoo = jsparse.BCOO.fromdense(mat)
vec_jax = jnp.array(vec)

def func(M, v):
  return 2 * (M @ v)

func_jax = jax.jit(func)

# trigger JIT compilation
_ = func_jax(mat_bcoo, vec_jax).block_until_ready()

print('jax operation')
%timeit func_jax(mat_bcoo, vec_jax).block_until_ready()
print()
print('scipy operation')
%timeit func(mat_coo, vec)

On a Colab CPU runtime I get this:

jax operation
6.25 ms ± 3.23 ms per loop (mean ± std. dev. of 7 runs, 100 loops each)

scipy operation
2.29 ms ± 294 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

While on a Colab GPU runtime I get this:

jax operation
211 µs ± 5.78 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

scipy operation
1.55 ms ± 94.4 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

As discussed in the FAQ entry I linked to, it's not surprising that for a small, self-contained operation like this run on CPU, the numpy ecosystem is more performant than JAX: this is exactly the domain that numpy/scipy operations have been optimized for! But you'll often find that JAX performance surpasses that of numpy-related tools when you're working on accelerators, and/or JIT-compiling longer programs.