RankedTensorType::get(static_cast<int32_t>(1), builder.getI32Type());
  auto out_segids_cst = builder.create<TFL::ConstOp>(
      builder.getI32TensorAttr(flattened_out_segids));
  auto contracting_segids_cst = builder.create<TFL::ConstOp>(
      builder.getI32TensorAttr(flattened_contracting_segids));
  auto num_segids_tensor =
      builder.create<TFL::ConstOp>(DenseIntElementsAttr::get(

           "Disable deadness related safety checks when clustering (this is "
           "unsound)."),
      Flag("tf_xla_disable_resource_variable_safety_checks_for_debugging",
           &mark_for_compilation_flags
                ->tf_xla_disable_resource_variable_safety_checks_for_debugging,
           "Disable resource variables related safety checks when clustering "
           "(this is unsound)."),

  explicit LiteSessionWrapper(std::unique_ptr<Session> wrapped)
      : wrapped_(std::move(wrapped)) {}

  Status Create(const GraphDef& graph) override {
    return absl::UnimplementedError("Session::Create()");
  }
  Status Create(GraphDef&& graph) override {
    return absl::UnimplementedError("Session::Create()");
  }

  Status Extend(const GraphDef& graph) override {
    return absl::UnimplementedError("Session::Extend()");

    return a->isBeforeInBlock(b);
  });

  // Create tf_device::ClusterOp before the last operation in the block that
  // is a part of a match set.
  auto back = cluster.operations.back();
  auto loc = back->getLoc();
  OpBuilder builder(back);

  auto cluster_op =
      builder.create<tf_device::ClusterOp>(loc, return_types, policy);

  // Create block in cluster_op's region and move 'cluster.operations' into

  SliceInputs slice_inputs_int64 =
      MakeSliceIndexAndSizeInt64(host_scope, slice_inputs);

  // Create a list of all control dependencies to be copied when possibly
  // replacing nodes related to slice_size.
  Node* old_size;
  std::vector<const Edge*> old_size_ctrl_deps;
  TF_RETURN_IF_ERROR(slice->input_node(2, &old_size));

  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();

  // 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;