#include <memory>

#include "mlir/Pass/Pass.h"  // from @llvm-project

namespace mlir {
namespace odml {

std::unique_ptr<Pass> CreateRenameEntrypointToMainPass();

}  // namespace odml
}  // namespace mlir

    }
    // Found entrypoint
    Builder builder(candidate_funcs[0]);
    candidate_funcs[0].setName(builder.getStringAttr("main"));
  }
};

std::unique_ptr<Pass> CreateRenameEntrypointToMainPass() {
  return std::make_unique<RenameEntrypointToMainPass>();
}

static PassRegistration<RenameEntrypointToMainPass> pass;

}  // namespace odml

namespace mlir {
namespace odml {

void AddTFToStablehloPasses(OpPassManager& pm, bool skip_resize,
                            bool smuggle_disallowed_ops) {
  pm.addPass(CreateRenameEntrypointToMainPass());

  // if the input is a call_xla_module, then unwrap the content
  pm.addPass(mlir::odml::CreateLegalizeTFXlaCallModuleToStablehloPass());

void AddTFToStablehloPasses(
    mlir::PassManager& pm,
    llvm::ArrayRef<llvm::ArrayRef<int64_t>> input_arg_shapes) {
  pm.addPass(mlir::odml::CreateRenameEntrypointToMainPass());
  // TODO: b/230572023 - Consider improving shape inference for While op instead
  // of dropping the attribute. This need not be correct for models not trained
  // on TPU.