code: def get_params_format_fn(
      init_params: FrozenDict,
  ) -> Tuple[int, Callable]:
      """Return a function that formats the parameters into a correct format."""
  
      flat, tree = tree_util.tree_flatten(init_params)
      params_sizes = np.cumsum([np.prod(p.shape) for p in flat])
  
      def params_format_fn(params: jnp.ndarray) -> FrozenDict:
          params = tree_util.tree_map(
              lambda x, y: x.reshape(y.shape),
              jnp.split(params, params_sizes, axis=-1)[:-1],
              flat,
          )
          return tree_util.tree_unflatten(tree, params)
  
      return params_sizes[-1], params_format_fn
  
file_path: evojax/evojax/util.py
imports: import os
  import logging
  import numpy
  from typing import Union
  from typing import Tuple
  from typing import Callable
  import jax.numpy
  from jax import tree_util
  from flax.core import FrozenDict
label_desc: Use image node_id 1 for calling a global function with example usage: get_params_format_fn(init_params) and returns: params_sizes,  params_format_fn, tree_util