code: def get_sparsity(params: hk.Params):
      """Calculate the total sparsity and tensor-wise sparsity of params."""
      total_params = sum(
          jnp.size(x) for x in jax.tree_util.tree_leaves(params)
      )
      total_nnz = sum(
          jnp.sum(x != 0.0) for x in jax.tree_util.tree_leaves(params)
      )
      leaf_sparsity = jax.tree_util.tree_map(
          lambda x: jnp.sum(x == 0) / jnp.size(x), params
      )
      return total_params, total_nnz, leaf_sparsity
  
imports: from collections.abc import Iterator, Mapping, Sequence
  import functools
  from typing import Callable
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 7 for calling a global function with example usage: get_sparsity(params) and returns: total_params,  total_nnz,  leaf_sparsity