code: def nested_kl_divergence(
      nested_from_distribution: types.NestedDistribution,
      nested_to_distribution: types.NestedDistribution,
      outer_dims: Sequence[int] = (),
  ) -> types.Tensor:
      """Given two nested distributions, sum the KL divergences of the leaves."""
      nest_utils.assert_same_structure(
          nested_from_distribution, nested_to_distribution
      )
  
      # Make list pairs of leaf distributions.
      flat_from_distribution = tf.nest.flatten(nested_from_distribution)
      flat_to_distribution = tf.nest.flatten(nested_to_distribution)
      all_kl_divergences = [
          from_dist.kl_divergence(to_dist)
          for from_dist, to_dist in zip(
              flat_from_distribution, flat_to_distribution
          )
      ]
  
      all_kl_divergences_reduced = []
      for kl_divergence in all_kl_divergences:
          # Reduce_sum over non-batch dimensions.
          reduce_dims = list(range(len(kl_divergence.shape)))
          for dim in outer_dims:
              reduce_dims.remove(dim)
          all_kl_divergences_reduced.append(
              tf.reduce_sum(
                  input_tensor=kl_divergence, axis=reduce_dims
              )
          )
  
      # Sum the kl of the leaves.
      total_kl = tf.add_n(all_kl_divergences_reduced)
  
      return total_kl
  
file_path: agents/tf_agents/agents/ppo/ppo_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Sequence
  import tensorflow
  import tensorflow_probability
  from tf_agents.distributions import utils
  from tf_agents.specs import distribution_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 5 for calling a global function with example usage: nested_kl_divergence(nested_from_distribution, nested_to_distribution, outer_dims) and returns: total_kl