bool skip_quantization_ops,
                                    bool skip_resize,
                                    bool skip_partitioned_calls) {
  TFToStablehloOptions options;
  options.skip_quantization_ops = skip_quantization_ops;
  options.skip_resize = skip_resize;
  options.skip_partitioned_calls = skip_partitioned_calls;
  PopulateLegalizeTFToStablehloPipeline(pm, options);
}

// Adds all the necessary passes to lower a TF module to StableHLO.
// `skip_resize` enables or disables skipping conversion of tf.ResizeBilinear
// and tf.ResizeNearestNeighbor ops.
// `smuggle_disallowed_ops` enables or disables converting disallowed ops
// like tf.ResizeBilinear or tf.ResizeNearestNeighbor to mhlo.custom_call ops.
void AddTFToStablehloPasses(OpPassManager& pm, bool skip_resize,
                            bool smuggle_disallowed_ops);

#include "tensorflow/compiler/mlir/tf2xla/transforms/passes.h"
#include "xla/mlir_hlo/mhlo/transforms/passes.h"

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

// Adds passes which transform TF Ops to StableHLO Ops.
void AddLegalizeTFToStablehloPasses(OpPassManager& pm,
                                    bool skip_quantization_ops,
                                    bool skip_resize,
                                    bool skip_partitioned_calls);

}  // namespace odml
}  // namespace mlir

opt<bool> skip_checks("skip-checks",
                      llvm::cl::desc("Skip checking for disallowed ops."),
                      llvm::cl::Optional, llvm::cl::init(false));

// NOLINTNEXTLINE
opt<bool> skip_resize(
    "skip-resize",
    llvm::cl::desc(
        "Skip converting tf.ResizeBilinear and tf.ResizeNearestNeighbor ops."),
    llvm::cl::Optional, llvm::cl::init(true));

// NOLINTNEXTLINE

  pm.addPass(mlir::TF::CreateTFShapeInferencePass());
  mlir::odml::AddLegalizeTFToStablehloPasses(pm, /*skip_quantization_ops=*/true,
                                             /*skip_resize=*/false,
                                             /*skip_partitioned_calls=*/false);
  pm.addNestedPass<mlir::func::FuncOp>(
      mlir::quant::stablehlo::createRemoveShardingCustomCallPass());

  // TF -> StableHLO legalization.
  // Skip StatefulPartitionedCall to preserve aliased functions.
  mlir::odml::AddLegalizeTFToStablehloPasses(pm, /*skip_quantization_ops=*/true,
                                             /*skip_resize=*/false,
                                             /*skip_partitioned_calls=*/true);
  // StableHLO -> MHLO legalization for MHLO optimization.
  pm.addPass(mlir::mhlo::createStablehloLegalizeToHloPass());

    }
    if (failed(pass_manager.run(module))) {
      return status_handler.ConsumeStatus();
    }
  }

  pass_manager.clear();
  mlir::odml::AddTFToStablehloPasses(pass_manager, /*skip_resize=*/true,
                                     /*smuggle_disallowed_ops=*/true);
  // Print out a detailed report of non-converted stats.
  pass_manager.addPass(mlir::odml::createPrintOpStatsPass(