code: def _sparse_bcoo_todense(
      state, jit: bool = False, compile: bool = False
  ):
      shape = (2000, 2000)
      nse = 10000
      size = math.prod(shape)
      rng = np.random.RandomState(1701)
      data = rng.randn(nse)
      indices = np.unravel_index(
          rng.choice(size, size=nse, replace=False), shape=shape
      )
      mat = sparse.BCOO(
          (jnp.array(data), jnp.column_stack(indices)), shape=shape
      )
  
      f = lambda mat: mat.todense()
      if jit or compile:
          f = jax.jit(f)
  
      if compile:
          while state:
              state.pause_timing()
              jax.clear_caches()
              state.resume_timing()
              f.lower(mat).compile()
      else:
          f(mat).block_until_ready()
          while state:
              f(mat).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 5 for calling a global function with example usage: _sparse_bcoo_todense(state, jit, compile) without return types