code: def _sparse_bcoo_matvec(
      state, jit: bool = False, compile: bool = False
  ):
      shape = (2000, 2000)
      nse = 10000
      key = jax.random.PRNGKey(1701)
      mat = sparse.random_bcoo(
          key,
          nse=nse,
          shape=shape,
          dtype=jnp.float32,
          indices_dtype=jnp.int32,
          sorted_indices=True,
      )
      vec = jax.random.uniform(
          key, shape=(shape[1],), dtype=jnp.float32
      )
  
      f = lambda mat, vec: mat @ vec
      if jit or compile:
          f = jax.jit(f)
  
      if compile:
          while state:
              state.pause_timing()
              jax.clear_caches()
              state.resume_timing()
              f.lower(mat, vec).compile()
      else:
          f(mat, vec).block_until_ready()
          while state:
              f(mat, vec).block_until_ready()
  
file_path: jax/benchmarks/sparse_benchmark.py
imports: from functools import partial
  import jax.numpy
  import numpy
  import math
  import google_benchmark
  import jax
  from jax.experimental import sparse
label_desc: Use image node_id 9 for calling a global function with example usage: _sparse_bcoo_matvec(state, jit, compile) without return types