code: def _pytree_unflatten(cls, aux_data, children):
      keys, metadata = aux_data
      model_dists = dict(list(zip(keys, children)))
      model = _unflatten_model(model_dists, metadata["callable_params"])
      return cls(model, **metadata["non_tensor_params"])
  
imports: import collections
  import functools
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import joint_distribution
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.internal import auto_composite_tensor
  from tensorflow_probability.python.internal import distribution_util
  from tensorflow.python.util import tf_inspect
label_desc: Use image node_id 7 for calling a global function with example usage: _pytree_unflatten(cls, aux_data, children) and returns: cls