Value tensor_i32 = builder.create<TF::CastOp>(
      loc, tensor_type.clone(builder.getIntegerType(32)), tensor);
  auto reduced =
      builder.create<TF::SumOp>(loc, tensor_i32, reduction_indices_value,
                                /*keep_dims=*/builder.getBoolAttr(true));
  auto mul_op = builder.create<TF::MulOp>(loc, zp, reduced);

    is_list_output.push_back(is_list);
  }

  strings::StrAppend(&comment, "\nReturns:\n");
  if (graph_op_def.output_arg_size() == 0) {  // No outputs.
    strings::StrAppend(&comment, "* the created `Operation`\n");
  } else if (graph_op_def.output_arg_size() == 1) {  // One output
    if (is_list_output[0]) {
      strings::StrAppend(&comment, "* `OutputList`: ");
    } else {

  // Add i32 -> i8 requantization.
  UniformQuantizeOp uniform_quant_op = rewriter.create<UniformQuantizeOp>(
      op.getLoc(), func_result_type, op.getResults());
  cast<func::ReturnOp>(entry_func_op.getBody().front().getTerminator())
      .setOperand(0, uniform_quant_op);
}

template <typename GemmStyleOp>
// Creates a quantized bias pattern for static and dynamic shape case

// be created).
static void DumpModule(mlir::ModuleOp module, std::string file_prefix) {
  std::string prefix = GetDumpDirFromEnvVar();
  if (prefix.empty()) return;

  auto* env = tensorflow::Env::Default();
  auto status = env->RecursivelyCreateDir(prefix);
  if (!status.ok()) {
    LOG(WARNING) << "cannot create directory '" << prefix
                 << "': " << status.message();

    // if a special Identity node in the following pattern is clustered in.
    // That is, an Identity node in the loop cond computation is used to drive
    // const nodes consumed by the loop body.  If this Identity node goes into
    // the same cluster with nodes from the loop body, extra dependency is
    // created between the loop cond and body computations and it hinders the

  // It's here so xla::Parameters that are created form block.getArguments will
  // have the proper shapes.
  TF_RETURN_IF_ERROR(RefineShapes(arg_shapes, module_op));

  mlir::func::FuncOp main = module_op.lookupSymbol<mlir::func::FuncOp>("main");
  mlir::Block& block = main.getRegion().front();
  xla::XlaBuilder builder("main");

  // Create xla_params.
  std::vector<xla::XlaOp> xla_params;

    // be categorized as cusom ops despite each of them may require different
    // behaviors. In that case, these ops can be marked in the custom map and
    // treated separately in this pass.

    auto custom_op = llvm::dyn_cast_or_null<CustomOp>(op);
    if (!custom_op) return false;

    // Custom op which is marked in the custom op map is quantizable.

    new_operands.push_back(block->addArgument(t, loc));
  auto call = builder.create<func::CallOp>(loc, func, new_operands);
  builder.create<YieldOp>(loc, call.getResults());
  // Mark old function as private so that it can be DCE'd if not called.
  func.setPrivate();
}

void RunOnWhile(TF::WhileOp while_op) {
  Operation* op = while_op.getOperation();
  // Create new TFL While op that will be used to replace TF While op.

  if (failed(applyPatternsAndFoldGreedily(module_op, std::move(patterns)))) {
    module_op.emitError() << "quant-cast-bf16-ops-to-f32 failed.";
    signalPassFailure();
  }
}

}  // namespace

// Creates an instance of the Cast BF16 ops to F32 pass.
std::unique_ptr<OperationPass<ModuleOp>> CreateCastBf16OpsToF32Pass() {
  return std::make_unique<CastBf16OpsToF32Pass>();
}

      "_xla", "_oc", "cluster", name_attrs, "cluster_rewritten", "host_graph",
      host_compute_core, &fld, &shape_inference_graphs,
      &has_outside_compilation));

  // Check host graph is not created.
  EXPECT_EQ(fld.Find("host_graph"), nullptr);
}

TEST_F(ExtractOutsideCompilationForFunctionTest, OutsideCompilationInIf) {
  // Build the XLA computation func.
  // "const0" (bool)
  // "const1" (int32)