def _tf_function_function_handler(input_signature):
  """Call function in eager mode, but also wrapped in `tf.function`."""
  def wrapper(fn):
    wrapped_fn = tf.function(fn, input_signature)
    return _eager_function_handler(input_signature)(wrapped_fn)
  return wrapper
Call function in eager mode, but also wrapped in `tf.function`.
tensorflow_transform/test_case.py