code: def _coo_extract(row: Array, col: Array, mat: Array) -> Array:
      """Extract values of dense matrix mat at given COO indices."""
      return mat[row, col]
  
imports: import functools
  from typing import Any, NamedTuple, Union
  import numpy
  import jax
  from jax import lax
  from jax import tree_util
  from jax import vmap
  from jax._src import core
  from jax._src import dtypes
  from jax._src import stages
  from jax._src.api_util import flatten_axes
  import jax.numpy
  from jax.util import safe_zip
  from jax._src.lax.lax import _dot_general_shape_rule, DotDimensionNumbers
  from jax._src.typing import Array
label_desc: Use image node_id 5 for calling a global function with example usage: _coo_extract(row, col, mat) and returns: mat