// Add inline pass.
  pass_manager.addPass(mlir::createInlinerPass());

  // Expands mhlo.tuple ops.
  pass_manager.addPass(
      mlir::mhlo::createExpandHloTuplesPass(entry_function_name.str()));
  // Flatten tuples for control flows.
  pass_manager.addNestedPass<mlir::func::FuncOp>(
      mlir::mhlo::createFlattenTuplePass());

  mlir::odml::AddMhloOptimizationPasses(
      pass_manager,