: input_shapes_(input_shapes) {}
  void runOnOperation() override {
    // Parse `input_arg_shapes_` if provided (test only)
    SmallVector<ArrayRef<int64_t>> input_shapes_vec;
    absl::StatusOr<SmallVector<SmallVector<int64_t>>> parsed_shapes;
    if (!input_arg_shapes_.empty()) {
      parsed_shapes = ParseArgumentShapes(input_arg_shapes_);
      if (!parsed_shapes.ok()) {

  let constructor = "TF::CreateTFShapeInferencePass()";

  let options = [
    Option<"max_iterations_", "max-iterations", "int64_t", /*default=*/"10",
           "Maximum shape inference iterations">,
    Option<"input_arg_shapes_", "input-arg-shapes", "std::string", /*default=*/"",
           "Input tensor shapes. Shapes for different tensors are separated by ':', and dimension sizes for the same tensor are separated by ','">,
  ];
}

      /*deserialize_xla_call_module=*/false, /*input_arg_shapes=*/{});
}

// TF->StableHLO has limited support for dynamic shapes.
// Some models can only be converted with explicitly provided input argument
// shapes.
void AddTFToStablehloPasses(
    mlir::PassManager& pm,
    llvm::ArrayRef<llvm::ArrayRef<int64_t>> input_arg_shapes = {});

}  // namespace quantization
}  // namespace tensorflow

// StableHLO module serialized into a XlaCallModuleOp. (ex: JAX/PyTorch models)
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

    session = saved_model_bundle->GetSession();
  }
  TF_ASSIGN_OR_RETURN(auto input_arg_shapes_vec,
                      TF::ParseArgumentShapes(input_arg_shapes_str));
  llvm::SmallVector<llvm::ArrayRef<int64_t>> input_arg_shapes(
      input_arg_shapes_vec.begin(), input_arg_shapes_vec.end());
  TF_RETURN_IF_ERROR(tensorflow::quantization::PreprocessAndFreezeGraph(
      /*mlir_dump_file_prefix=*/"", /*is_inliner_run=*/true,