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,

      auto body_yield_operand =
          LookThroughIdentity(body_yield.getOperand(op_idx));
      auto cond_yield_operand =
          cond_forwards_args
              ? LookThroughIdentity(cond_yield.getOperand(op_idx + 1))
              : nullptr;
      auto while_operand = while_op.getOperand(op_idx);
      if ((body_arg == body_yield_operand ||
           while_operand == body_yield_operand) &&

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

  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,

    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;

  // 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 mlir::isa<mlir::TF::StringType>(type.getElementType());
  });
  return has_string;
}

bool HasStringOperand(Operation& op) {
  for (auto operand : op.getOperands()) {
    auto operand_type = getElementTypeOrSelf(operand);
    if (IsStringType(operand_type)) return true;
  }
  return false;
}

bool HasStringResult(Operation& op) {

              1, mlir::stablehlo::PrecisionAttr::get(
                     ctx_.get(), mlir::stablehlo::Precision::DEFAULT)))),
  };
  const SmallVector<Value> operands(dot_general_op->getOperands());
  const SmallVector<Value> results(dot_general_op->getResults());
  Operation* lifted_op =
      LiftAsFunctionCall(builder_, dot_general_op->getLoc(),
                         FunctionCallOpType::TFXlaCallModuleOp,

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

      for (auto result : filter_resources(op->getResults())) {
        for (auto operand : op->getOperands()) {
          PropagateInputToOutput(operand, result);
        }
      }
    } else if (llvm::isa<IdentityNOp, IdentityOp, CastOp>(op)) {
      for (auto result : filter_resources(op->getResults()))
        PropagateInputToOutput(op->getOperand(result.getResultNumber()),
                               result);