mhlo::ReturnOp return_op = dyn_cast<mhlo::ReturnOp>(body.back());
  if (!return_op) return failure();
  if (return_op.getNumOperands() != 1) return failure();
  if (return_op.getOperands().front() != body.getArgument(1)) return failure();
  return success();
}

bool NeedsReformatTypeAndPermutation(int batch_dim, int feature_dim,
                                     int spatial_dim_start,

  llvm::SetVector<Value> been_defined;
  llvm::SetVector<Value> results;
  auto update_from_op = [&](Operation* op) {
    been_defined.insert(op->getResults().begin(), op->getResults().end());
    for (Value input : op->getOperands()) {
      if (been_defined.contains(input)) {
        continue;
      }
      results.insert(input);
    }
  };
  for (Operation* op : partition_ops) {
    update_from_op(op);

        ReadAttr(flex_buffer_map, &builder, custom_op.getCustomCode().str());
    OperationState op_state(custom_op.getLoc(),
                            custom_op.getCustomCode().str());
    op_state.addOperands(custom_op.getOperands());
    llvm::SmallVector<mlir::Type, 4> output_tys;
    for (int i = 0; i < custom_op.getNumResults(); i++) {
      output_tys.push_back(custom_op.getType(i));
    }
    op_state.addTypes(output_tys);

      // inputs.
      SmallVector<Value, 4> inputs;
      inputs.reserve(quantizing_op->getNumOperands());
      for (const auto& operand : quantizing_op->getOperands()) {
        Type operand_type = operand.getType();
        if (operand_type.isa<NoneType>()) {
          inputs.push_back(operand);
          continue;
        }

    }
  }
  return false;
}

bool HasDynamicExternalValues(Operation* op) {
  return op
      ->walk([](Operation* walked_op) {
        for (Value v : walked_op->getOperands()) {
          if (mlir::TF::CanBeRefined(v.getType())) {
            return WalkResult::interrupt();
          }
        }
        return WalkResult::advance();
      })
      .wasInterrupted();
}

    const std::string tf_op_full_name = llvm::Twine("tf.", tf_op_name).str();

    // Create the TF op
    OperationState op_state(op.getLoc(), tf_op_full_name);
    op_state.addOperands(op.getOperands());
    op_state.addTypes(op.getResultTypes());

    SmallVector<NamedAttribute, 2> attrs;
    std::string parsed_op_name;
    tensorflow::NodeDef node_def;

  llvm::SetVector<Operation*> new_ops;

  while (!ops_to_process.empty()) {
    for (Operation* op : ops_to_process) {
      if (predecessors) {
        for (Value operand : op->getOperands()) {
          // Stop at the block boundary.
          if (mlir::isa<BlockArgument>(operand)) continue;

          Operation* predecessor = operand.getDefiningOp();
          if (!operations->contains(predecessor) &&

  // then, they are represented as mlir::TF::IfOp. We add them here, anyway, to
  // be future-proof.
  if (llvm::isa<TF::IfOp, TFL::IfOp, TFL::CallOnceOp, TFL::WhileOp>(op))
    return true;
  for (auto operand : op->getOperands()) {
    if (IsResourceTensor(operand)) return true;
  }
  for (auto result : op->getResults()) {
    if (IsResourceTensor(result)) return true;
  }
  return false;
}

  return wrapped_op.getNumResults() == terminator->getNumOperands() &&
         std::equal(wrapped_op.getResults().begin(),
                    wrapped_op.getResults().end(),
                    terminator->getOperands().begin());
}
}  // end anonymous namespace

TensorFlowDeviceDialect::TensorFlowDeviceDialect(MLIRContext* context)
    : Dialect(/*name=*/"tf_device", context,

    }
  }
  OpBuilder builder(op);
  builder.setInsertionPointAfter(op);
  OperationState state(op->getLoc(), op->getName().getStringRef(),
                       op->getOperands(), result_types, op->getAttrs());
  for (int i = 0; i < op->getNumRegions(); ++i) {
    state.addRegion();
  }
  Operation* new_op = builder.create(state);