code: def _dynamic_partition_of_nested_tensors(
      nested_tensor: types.NestedTensor,
      partitions: types.Int,
      num_partitions: int,
  ) -> List[types.NestedTensor]:
      """This function takes a nested structure and partitions every element of it.
  
      Specifically it outputs a list of nest that all have the same structure as the
      original, and every element of the list is a nest that contains a dynamic
      partition of the corresponding original tensors.
  
      Note that this function uses tf.dynamic_partition, and thus
      'MixtureAgent' is not compatible with XLA.
  
      Args:
        nested_tensor: The input nested structure to partition.
        partitions: int32 tensor based on which the partitioning happens.
        num_partitions: The number of expected partitions.
  
      Returns:
        A list of nested tensors with the same structure as `nested_tensor`.
      """
      flattened_tensors = tf.nest.flatten(nested_tensor)
      if not flattened_tensors:
          return [nested_tensor] * num_partitions
      partitioned_flat_tensors = [
          tf.dynamic_partition(
              data=t,
              partitions=partitions,
              num_partitions=num_partitions,
          )
          for t in flattened_tensors
      ]
      list_of_partitions = list(
          map(list, zip(*partitioned_flat_tensors))
      )
      return [
          tf.nest.pack_sequence_as(nested_tensor, i)
          for i in list_of_partitions
      ]
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  from typing import List, Optional, Sequence, Text
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.policies import mixture_policy
  from tf_agents.trajectories import trajectory
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 for calling a global function with example usage: _dynamic_partition_of_nested_tensors(nested_tensor, partitions, num_partitions) and returns: unknown, unknown