void RemoveClusterAliasedOutputs(OpBuilder* builder,
                                 mlir::tf_device::ClusterOp cluster) {
  llvm::SmallVector<Value, 4> used_old_cluster_results;
  llvm::SmallVector<Value, 4> new_cluster_results;
  llvm::SmallVector<Type, 4> new_cluster_result_types;
  Operation* cluster_terminator = cluster.GetBody().getTerminator();
  for (auto result :

void FindCustomAggregatorOps(
    Region& region,
    const std::unordered_set<std::string>& aggregator_ops_to_ignore,
    SmallVector<Value>& statistics_outputs, SmallVector<StringRef>& ids,
    SmallVector<int32_t>& calibration_methods) {
  for (auto op : region.getOps<TF::CustomAggregatorOp>()) {
    if (aggregator_ops_to_ignore.count(op.getId().str())) continue;

#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_COMMON_LIFT_AS_FUNCTION_CALL_H_

#include "absl/base/nullability.h"
#include "absl/status/statusor.h"
#include "llvm/ADT/SmallVector.h"
#include "mlir/Bytecode/BytecodeOpInterface.h"  // from @llvm-project
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project

// Move Transpose operations that permute `op` operands after the `op`.
void MoveTransposeAfter(Operation* op, SmallVector<Operation*, 8>* work_list,
                        bool fold_transpose_in_ops) {
  // Indices of operands and results that depend on data layout.
  SmallVector<unsigned, 4> layout_dependent_operands;
  SmallVector<unsigned, 4> layout_dependent_results;

  auto fold_operands = dyn_cast<FoldOperandsTransposeInterface>(op);

  // TF::UnsortedSegmentProdOp.
  llvm::SmallVector<int32_t, 4> flattened_out_segids =
      llvm::SmallVector<int32_t, 4>(operand_rank, static_cast<int32_t>(-1));
  for (int64_t i : dot_dimensions_info.out_dimensions().AxesArray()) {
    flattened_out_segids[i] = 0;
  }
  llvm::SmallVector<int32_t, 4> flattened_contracting_segids =
      llvm::SmallVector<int32_t, 4>(operand_rank, static_cast<int32_t>(-1));

    mlir::tf_device::ReplicateOp replicate_op,
    PartitionedClusterOutputMap& partitioned_outputs,
    llvm::SmallVector<mlir::TF::IdentityOp, 8>& erase_list,
    llvm::SmallVector<Type, 8>& result_types, int num_replicas) {
  Operation* result_op;
  llvm::SmallVector<Value, 8> results;
  uint64_t num_results = cluster.getNumResults();
  for (uint64_t result_id = 0; result_id < num_results; ++result_id) {

limitations under the License.
==============================================================================*/

// This pass remove Cluster ops by inlining Cluster ops.

#include "llvm/ADT/SmallVector.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_device.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

          if (func && reachable_funcs.insert(func).second)
            funcs_to_visit.push_back(func);
        });
      };

  SmallVector<func::FuncOp, 4> funcs_to_visit;
  collect_reachable_funcs(region, funcs_to_visit);

  while (!funcs_to_visit.empty()) {
    SmallVector<func::FuncOp, 4> new_funcs_to_visit;
    for (func::FuncOp func_to_visit : funcs_to_visit) {
      if (!func_to_visit.getCallableRegion()) continue;

    ArrayRef<int64_t> transpose_perm = transpose_op.getPermutation();
    SmallVector<int64_t> new_padding_low =
        InvertPermutationToVector(pad_op.getEdgePaddingLow(), transpose_perm);
    SmallVector<int64_t> new_padding_high =
        InvertPermutationToVector(pad_op.getEdgePaddingHigh(), transpose_perm);
    SmallVector<int64_t> new_padding_interrier =

  // Pack op should have same shape type.
  SmallVector<Value, 5> pack_inputs(pack_op.getValues());
  auto input_type = mlir::dyn_cast<RankedTensorType>(pack_inputs[0].getType());
  if (!input_type) return failure();

  // Figure out output shapes.
  SmallVector<int64_t, 4> concat_out_shape;
  SmallVector<int64_t, 4> pack_out_shape;

  const int64_t rank = input_type.getRank();