def decorator(fn):
    """Logs function calls in a tensorflow graph collection."""

    @functools.wraps(fn)
    def wrapped_fn(*args, **kwargs):
      if not _in_logging_context:
        with logging_context():
          graph = tf.compat.v1.get_default_graph()
          # Collection is a list that contains a single Counter of {name: count}
          # Note: We aggregate counts of function calls instead having one
          # collection item per call, since TFT users can use an arbitrarily
          # large number of analyzers and mappers and we don't want the graph
          # to get too big.
          # TODO(rachelim): Make this collection serializable so it can be added
          # to the SavedModel.
          collection = graph.get_collection_ref(collection_name)
          if not collection:
            collection.append(collections.Counter())
          collection[0][fn.__name__] += 1
          return fn(*args, **kwargs)
      else:
        return fn(*args, **kwargs)

    # We use tf_decorator here so that TF can correctly introspect into
    # functions for docstring generation.
    return tf_decorator.make_decorator(fn, wrapped_fn)
Logs function calls in a tensorflow graph collection.
tensorflow_transform/common.py