device_cluster.getLoc(), num_regions, llvm::ArrayRef<Type>{});
  SmallVector<Operation*, 4> core_to_host_insertion_point;
  SmallVector<mlir::tf_device::LaunchOp, 4> core_to_tmp_launch;
  SmallVector<Operation*, 4> compilation_key_ops;
  SmallVector<Value, 4> core_to_compilation_key;
  SmallVector<Operation*, 4> core_to_device_ordinal_op;
  SmallVector<Value, 4> core_to_device_ordinal;

#include "mlir/Support/LLVM.h"  // from @llvm-project

namespace mlir {
namespace TF {

SmallVector<int64_t, 4> ReversePermutation(ArrayRef<int64_t> permutation) {
  SmallVector<int64_t, 4> reverse(permutation.size());
  for (size_t i = 0; i < permutation.size(); ++i) {
    reverse[permutation[i]] = i;
  }
  return reverse;
}

SmallVector<int64_t, 4> GetDataFormatPermutation(StringRef from, StringRef to) {

      for (Operation& op : nested_block->getOperations()) update_from_op(&op);
    });
  }
  return SmallVector<Value>(results.getArrayRef());
}

llvm::SmallVector<Value> AccumulateResultsDefinedWithin(
    const llvm::SetVector<Operation*>& partition_ops) {
  llvm::SmallVector<Value> values_for_results;
  for (Operation* op : partition_ops) {
    if (IsConstantOrNone(op)) {
      continue;
    }

func::FuncOp createLstmCompositeFunc(mlir::Builder* builder, bool ln,
                                     bool cifg) {
  SmallVector<int64_t, 2> input_shape{1, 2};
  SmallVector<int64_t, 2> weight_shape{3, 12};
  SmallVector<int64_t, 1> bias_shape{2};
  SmallVector<int64_t, 2> projection_shape{1, 2};
  SmallVector<int64_t, 1> layer_norm_scale{4};
  SmallVector<int64_t, 2> output_shape{1, 2};
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());

namespace mlir {
namespace quant {
namespace {

// Folds the operation recursively and return the results.
LogicalResult FoldOperation(OpBuilder& builder, Operation* op,
                            SmallVector<Value>& results) {
  SmallVector<ElementsAttr> inputs;
  for (auto operand : op->getOperands()) {
    auto preceding_const_op = operand.getDefiningOp<TF::ConstOp>();
    if (preceding_const_op) {

  assert(devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))
             ->getSecond()
             .size() == num_replicas);

  llvm::SmallVector<std::pair<ValueRange, Type>, 8> new_replicated_inputs;
  llvm::SmallVector<Value, 8> new_packed_inputs;
  llvm::SmallVector<llvm::SmallVector<Value, 8>, 8> replicated_inputs;
  replicated_inputs.reserve(replicate.GetNumReplicatedBlockArguments());

  builder->create<tf_device::ReturnOp>(loc, op->getResults());

  // Move op inside cluster.
  op->moveBefore(launch.GetBody().getTerminator());
}

llvm::SmallVector<Operation*, 4> GetResourceOpsUsingCompositeArgsInReplicate(
    tf_device::ReplicateOp replicate) {
  llvm::SmallVector<Operation*, 4> resource_users;
  const auto add_resource_op_to_list = [&resource_users](Operation* op) {

// `tf.AssignVariableOp`s and `tf.Const`s. Collects and returns the
// `tf.VarHandleOp`s that are initialized by these `tf.AssignVariableOp`s.
SmallVector<TF::VarHandleOp> CollectVariableOps(
    func::FuncOp session_init_func) {
  SmallVector<TF::VarHandleOp> var_handle_ops{};

  for (auto assign_variable_op : llvm::make_early_inc_range(
           session_init_func.getOps<TF::AssignVariableOp>())) {

AnonymousIteratorV3Op CreateIterator(OpBuilder builder,
                                     llvm::ArrayRef<Type> dataset_types,
                                     ReduceDatasetOp reduce_dataset) {
  llvm::SmallVector<Attribute, 2> shape_attrs;
  llvm::SmallVector<Attribute, 2> type_attrs;
  for (Type type : dataset_types) {
    shape_attrs.push_back(
        TF::ShapeAttr::get(builder.getContext(), mlir::cast<ShapedType>(type)));

};

// Compute the permutation that would take `arr` to the identity.
llvm::SmallVector<int32_t> InvertPermutation(llvm::SmallVector<int32_t> arr) {
  llvm::SmallVector<int32_t> inverse_arr(arr.size());
  for (int32_t i = 0; i < arr.size(); ++i) {
    inverse_arr[arr[i]] = i;
  }
  return inverse_arr;
}

llvm::SmallVector<int64_t> PermuteShape(llvm::ArrayRef<int64_t> shape,