def write_v2_saved_model(
    module: tf.Module,
    tf_function: tf.types.experimental.GenericFunction,
    name: str,
    saved_model_dir: str,
    save_options: Optional[tf.saved_model.SaveOptions] = None,
) -> function.ConcreteFunction:
  """Writes `tf_function` under attr `name` of `module` to `saved_model_dir`."""
  concrete_fn = trace_and_update_module(
      module, tf_function, name, strip_control_dependencies=False
  )
  tf.saved_model.save(
      module,
      saved_model_dir,
      options=save_options,
  )
  return concrete_fn
Writes `tf_function` under attr `name` of `module` to `saved_model_dir`.
tensorflow_transform/saved/saved_transform_io_v2.py