builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, const llvm::SetVector<Operation*>& ops,

  for (auto& func : funcs_to_rewrite) {
    if (func.getSecond().original.isPublic() &&
        !func.getSecond().original.symbolKnownUseEmpty(module)) {
      auto clone = func.getSecond().original.clone();
      clone.setPrivate();
      symbol_table.insert(clone);
      func.getSecond().clone = clone;
    }
  }

  return success();
}

// Assigns either MAXIMAL or MANUAL sharding. The MAXIMAL sharding sends/recvs

  };
  func::FuncOp lowered_callee = callee;
  if (!callee.isPrivate()) {
    // Clone non-private callee in case of signature change.
    lowered_callee = callee.clone();
    lowered_callee.setPrivate();
  }
  auto grads = AccessedGradients({lowered_callee}, module);
  for (int64_t i = 0; i < lowered_callee.getNumArguments(); ++i) {
    auto it = grads.find(i);
    if (it == grads.end()) continue;

        builder.getFunctionType(/*inputs=*/{}, /*results=*/{});
    auto summary_func = builder.create<func::FuncOp>(
        builder.getUnknownLoc(), /*sym_name=*/"summary", func_type);
    summary_func.setPrivate();
    summary_func->setAttr("quantization_summary",
                          builder.getStringAttr(log_message));
  }

 private:
  // Structs used to count quantized and non-quantized ops.

          subgraph, &builder, "outputs", func_outputs));
    }
    if (!attributes.empty()) {
      func->setAttr("tf.entry_function", builder.getDictionaryAttr(attributes));
    }
  } else {
    func.setPrivate();
  }

  // Set signature on function.
  if (signature) {
    SetSignature(func, signature, subgraph.tensors);
  }

  absl::flat_hash_set<const tflite::OperatorT*> pruned_subgraph_ops;

  SmallString<64> name_base = callee.getName();
  auto module = callee->getParentOfType<ModuleOp>();
  name_base += "_resource_lifted";
  auto name = name_base;
  callee = callee.clone();
  callee.setPrivate();
  callee.setName(mlir::StringAttr::get(callee->getContext(), name));
  SymbolTable(module).insert(callee);
  result->lifted_callee = callee;

  // Remove unused resources in functions.