MLIRContext* context = backprop.getContext();
  SmallVector<int64_t, 4> values = {1, 1, 1, 1};
  auto attrs = llvm::map_range(values, [context](int64_t v) -> Attribute {
    return IntegerAttr::get(IntegerType::get(context, 64), APInt(64, v));
  });
  auto strides = ArrayAttr::get(context, llvm::to_vector<4>(attrs));

  // new result type.
  SmallVector<int64_t, 4> new_shape(new_filter_shape.begin(),

 private:
  // Collects all candidate ops for quantization, which are the
  // `dequantize_op`'s users.
  FailureOr<SmallVector<Operation*>> CollectCandidateOps(
      DequantizeOpT dequantize_op) const {
    auto users = dequantize_op->getResult(0).getUsers();
    return SmallVector<Operation*>(users.begin(), users.end());
  }

  // Collects all candidate ops for quantization, which is the operand of

  SmallVector<NamedAttribute> attributes;
  return LiftAsFunctionCall(builder, location, call_op_type, func_name,
                            arguments, results, attributes);
}

SmallVector<Value> AppendToVector(const ArrayRef<Value> arguments,
                                  Value append) {
  SmallVector<Value> ret(arguments);
  ret.push_back(append);
  return ret;
}

                           const SetVector<Operation*>& ops_to_add,
                           SmallVector<Operation*>& all_descendants) {
  if (!op) {
    return false;
  }

  SmallVector<Operation*> current_layer_descendants;
  SmallVector<Operation*> next_layer_descendants;
  int current_depth = 0;
  current_layer_descendants.push_back(op);
  // BFS downstream ops for current user.

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 :

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

    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) {

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

  ASSERT_TRUE(module);
  mlir::SymbolTable symtab(*module);
  llvm::SmallVector<mlir::func::FuncOp> entry_funcs =
      GetEntryFunctions(*module);
  EXPECT_EQ(entry_funcs.size(), 1);
  EXPECT_EQ(entry_funcs[0].getSymName(), "entry_func");
  llvm::SmallVector<mlir::SymbolUserOpInterface> outermost_pcall_ops;
  auto result =
      mlir::GetFirstOpsOfType<mlir::TF::StatefulPartitionedCallOp,

  const int operands_size = while_op.getNumOperands() + extra_operands.size();
  SmallVector<Value, 4> operands;
  operands.reserve(operands_size);
  operands.append(while_op.getOperands().begin(), while_op.getOperands().end());
  operands.append(extra_operands.begin(), extra_operands.end());
  SmallVector<Type, 4> new_types;
  new_types.reserve(operands_size);
  new_types.append(while_op.getResultTypes().begin(),