code: def _count_stored_elements_per_batch(
      mat: Array, n_batch: int = 0, n_dense: int = 0
  ) -> Array:
      """Return per-batch number of stored elements (nse) of a dense matrix."""
      mat = jnp.asarray(mat)
      mask = mat != 0
      if n_dense > 0:
          mask = mask.any(tuple(-(i + 1) for i in range(n_dense)))
      mask = mask.sum(tuple(range(n_batch, mask.ndim)))
      return mask
  
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 6 for calling a global function with example usage: _count_stored_elements_per_batch(mat, n_batch, n_dense) and returns: mask