return module_op;
  }

  // Evaluate return value of a function using TF kernel.
  // This assumes that the module op has only 1 function and it has TF ops only.
  absl::StatusOr<std::shared_ptr<xla::Literal>> EvaluateTfFunction(
      absl::string_view program,
      absl::Span<const xla::Literal* const> arguments) {
    TF_ASSIGN_OR_RETURN(auto module_op,
                        ReplaceFuncArgsByConstant(program, arguments));