namespace tensorflow {

// Converts the given Jax model to a TF Lite FlatBuffer
// string according to the given model flags, toco flags and tags. Returns error
// status if it fails to convert the input.
Status ConvertJaxToTFLiteFlatBuffer(const std::string& input,
                                    const toco::ModelFlags& model_flags,
                                    toco::TocoFlags& toco_flags,
                                    string* result);

  if (!status.ok()) {
    LOG(ERROR) << "HLO Module import failed: " << status;
    return nullptr;
  }

  return module;
}

}  // namespace
absl::Status ConvertJaxToTFLiteFlatBuffer(const std::string& input,
                                          const toco::ModelFlags& model_flags,
                                          toco::TocoFlags& toco_flags,

                      "Cannot specify both saved_model and hlo import.");
      return nullptr;
    }

    if (model_flags.use_hlo_import()) {
      status = tensorflow::ConvertJaxToTFLiteFlatBuffer(
          input_contents_txt, model_flags, toco_flags,
          &output_file_contents_txt);
    } else if (!model_flags.saved_model_dir().empty()) {
      status = tensorflow::ConvertSavedModelToTFLiteFlatBuffer(