code: def _function_with_flat_signature(
      function,
      input_specs,
      output_spec,
      include_batch_dimension,
      batch_size=None,
  ):
      """Create a tf.function with a given signature for export.
  
      Args:
        function: A callable that can be wrapped in tf.function.
        input_specs: A tuple nested specs declaring ordered arguments to function.
        output_spec: The nested spec describing the output of the function.
        include_batch_dimension: Python bool, whether to prepend a batch dimension
          to inputs and outputs.
        batch_size: Known batch size, or `None` for unknown.  Ignored if
          `include_batch_dimension == False`.
  
      Returns:
        A `tf.function` with the given input spec that returns a `dict` mapping
        output spec keys to corresponding output values.
      """
  
      def _with_batch(spec):
          if include_batch_dimension:
              return tf.TensorSpec(
                  shape=tf.TensorShape([batch_size]).concatenate(
                      spec.shape
                  ),
                  name=spec.name,
                  dtype=spec.dtype,
              )
          else:
              return spec
  
      flat_input_spec = [
          _with_batch(spec) for spec in tf.nest.flatten(input_specs)
      ]
  
      def as_dict(outputs, output_spec):
          nest_utils.assert_same_structure(outputs, output_spec)
          flat_outputs = tf.nest.flatten(outputs)
          flat_names = [s.name for s in tf.nest.flatten(output_spec)]
          return dict(zip(flat_names, flat_outputs))
  
      @common.function(input_signature=flat_input_spec)
      def function_with_signature(*input_list):
          inputs_ = tf.nest.pack_sequence_as(input_specs, input_list)
          outputs_ = function(*inputs_)
          dict_outputs_ = as_dict(outputs_, output_spec)
          return dict_outputs_
  
      return function_with_signature
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import functools
  import os
  from typing import Any, Callable, Dict, Tuple, Optional, Text, cast, Sequence
  from absl import logging
  import tensorflow
  import tensorflow_probability
  from tf_agents.policies import tf_policy
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
  from tensorflow.python.eager import def_function
label_desc: Use image node_id 6 for calling a global function with example usage: _function_with_flat_signature(function, input_specs, output_spec, include_batch_dimension, batch_size) and returns: function_with_signature, dict, dict_outputs_, tf, spec