code: def _add_tfma_collections(
      features: types.TensorTypeMaybeDict,
      labels: Optional[types.TensorTypeMaybeDict],
      input_refs: types.TensorType,
  ):
      """Add extra collections for features, labels, input_refs, version.
  
      This should be called within the Graph that will be saved. Typical usage
      would be when features and labels have been parsed, i.e. in the
      input_receiver_fn.
  
      Args:
        features: dict of strings to tensors representing features
        labels: dict of strings to tensors or a single tensor
        input_refs: See EvalInputReceiver().
      """
      # Clear existing collections first, in case the EvalInputReceiver was called
      # multiple times.
      del tf.compat.v1.get_collection_ref(
          encoding.with_suffix(
              encoding.FEATURES_COLLECTION, encoding.KEY_SUFFIX
          )
      )[:]
      del tf.compat.v1.get_collection_ref(
          encoding.with_suffix(
              encoding.FEATURES_COLLECTION, encoding.NODE_SUFFIX
          )
      )[:]
      del tf.compat.v1.get_collection_ref(
          encoding.with_suffix(
              encoding.LABELS_COLLECTION, encoding.KEY_SUFFIX
          )
      )[:]
      del tf.compat.v1.get_collection_ref(
          encoding.with_suffix(
              encoding.LABELS_COLLECTION, encoding.NODE_SUFFIX
          )
      )[:]
      del tf.compat.v1.get_collection_ref(
          encoding.EXAMPLE_REF_COLLECTION
      )[:]
      del tf.compat.v1.get_collection_ref(
          encoding.TFMA_VERSION_COLLECTION
      )[:]
  
      for feature_name, feature_node in features.items():
          _encode_and_add_to_node_collection(
              encoding.FEATURES_COLLECTION, feature_name, feature_node
          )
  
      if labels is not None:
          # Labels can either be a Tensor, or a dict of Tensors.
          if not isinstance(labels, dict):
              labels = {
                  util.default_dict_key(constants.LABELS_NAME): labels
              }
  
          for label_key, label_node in labels.items():
              _encode_and_add_to_node_collection(
                  encoding.LABELS_COLLECTION, label_key, label_node
              )
      # Previously input_refs was called example_ref. This code is being deprecated
      # so it was not renamed.
      example_ref_collection = tf.compat.v1.get_collection_ref(
          encoding.EXAMPLE_REF_COLLECTION
      )
      example_ref_collection.append(
          encoding.encode_tensor_node(input_refs)
      )
  
      tf.compat.v1.add_to_collection(
          encoding.TFMA_VERSION_COLLECTION, version.VERSION
      )
  
file_path: model-analysis/tensorflow_model_analysis/eval_saved_model/export.py
imports: from typing import Any, Callable, Dict, Optional, Union
  import tensorflow
  from tensorflow import estimator
  from tensorflow_model_analysis import version
  from tensorflow_model_analysis.api import types
  from tensorflow_model_analysis.eval_saved_model import constants
  from tensorflow_model_analysis.eval_saved_model import encoding
  from tensorflow_model_analysis.eval_saved_model import util
  from tensorflow_model_analysis.utils import util
  from tensorflow.python.estimator.export import export
label_desc: Use image node_id 3 for calling a global function with example usage: _add_tfma_collections(features, labels, input_refs) without return types