code: def log_param_shapes(
      params: Any,
      print_params_nested_dict: bool = False,
      description: Optional[str] = None,
  ) -> int:
      """Prints out shape of parameters and total number of trainable parameters.
  
      Args:
        params: PyTree of model parameters.
        print_params_nested_dict: If True, it prints parameters in shape of a nested
          dict.
        description: Optional description to print out before logging the parameter
          summary.
  
      Returns:
        int; Total number of trainable parameters.
      """
      if print_params_nested_dict:
          shape_dict = tree_map(lambda x: str(x.shape), params)
          # We use json.dumps for pretty printing nested dicts.
          logging.info(
              "Printing model param shape:/n%s",
              json.dumps(shape_dict, sort_keys=True, indent=4),
          )
      parameter_overview.log_parameter_overview(params, msg=description)
      total_params = jax.tree_util.tree_reduce(
          operator.add, tree_map(lambda x: x.size, params)
      )
      logging.info("Total params: %d", total_params)
      return total_params
  
imports: from collections import abc
  from concurrent import futures
  import json
  import operator
  import threading
  from typing import Any, Callable, Optional, Sequence, Set, Tuple, Union
  from absl import logging
  from clu import parameter_overview
  import jax
  import jax.numpy
  from jax.tree_util import tree_map
  import ml_collections
label_desc: Use image node_id 2 for calling a global function with example usage: log_param_shapes(params, print_params_nested_dict, description) and returns: total_params