// while creating the ops so make sure there aren't any type changes between
    // the original op operands and the operands during the conversion.
    for (auto&& [old_val, new_val] : llvm::zip(op->getOperands(), operands)) {
      if (old_val.getType() != new_val.getType()) return failure();
    }

    auto abstractOp = op->getRegisteredInfo();
    if (!abstractOp) return failure();

  for (auto *op : custom_ops) {
    builder.setInsertionPoint(op);
    Location loc = op->getLoc();
    auto custom_op = builder.create<CustomTfOp>(loc, op->getResultTypes(),
                                                op->getOperands());
    Region region;
    Block *new_block = new Block;
    region.push_back(new_block);

    builder.setInsertionPointToEnd(&region.front());
    Operation *inner_op = builder.clone(*op);

    // and mark the module as not valid for TFLite legalization.
    // Note: this might disable native variables in more than needed cases.
    // TODO(b/189370197): Enhance variable analysis.
    for (auto operand : op->getOperands()) {
      if (mlir::isa<TF::ResourceType>(
              getElementTypeOrSelf(operand.getType()))) {
        legalize_to_tfl = false;
        return WalkResult::interrupt();
      }
    }

  Operation* op = while_op.getOperation();
  // Create new TFL While op that will be used to replace TF While op.
  auto new_op = OpBuilder(op).create<TFL::WhileOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(),
      while_op.getIsStateless());
  Location loc = while_op->getLoc();
  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);
  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

    llvm::SetVector<Operation*>* unprocessed_subgraphs) {
  Subgraph& current_subgraph = subgraphs_.find(current_graph_call)->second;

  std::vector<Subgraph*> input_subgraphs;
  for (auto input : current_graph_call.getOperands()) {
    func::CallOp input_call = GetProducerCallOpOrNull(input);
    // If the input subgraph is not processed yet, we just go ahead and process
    // that one first.
    if (input_call == nullptr) continue;

    if (IsTrivialUnaryOperation(&cluster_op) ||
        HasOutsideCompilationAttribute(&cluster_op)) {
      auto walk_result = cluster_op.walk([&](Operation* op) {
        for (Value operand : op->getOperands()) {
          Operation* operand_op = GetOpOfValue(operand);
          if (head_outside_compiled_ops.count(operand_op)) continue;

          if (operand_op->getParentRegion() == cluster_region)

    // for function arguments.
    func.walk([&](Operation* op) {
      if (llvm::isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
        op->replaceAllUsesWith(op->getOperands());
        return;
      }
      if (auto grad = llvm::dyn_cast<TF::TensorArrayGradV3Op>(op)) {
        insert(grad.getHandle(), grad.getSource().str(), func_block);

  LogicalResult match(Operation* op) const override {
    if (isa<TF::CastOp, TF::ConstOp>(op) ||
        op->getName().hasTrait<OpTrait::ZeroOperands>()) {
      return failure();
    }
    for (Value input : op->getOperands()) {
      if (getElementTypeOrSelf(input).isBF16()) {
        return success();
      }
    }
    for (Value value : op->getResults()) {
      if (getElementTypeOrSelf(value).isBF16()) {
        return success();

// results.
void RemoveIdentityFromRegion(Region& region) {
  region.walk([](Operation* op) {
    if (isa<TF::IdentityOp, TF::IdentityNOp>(op)) {
      op->replaceAllUsesWith(op->getOperands());
      op->erase();
    }
  });
}

void TPUIdentityPruning::runOnOperation() {
  SmallVector<tf_device::ClusterOp, 4> clusters;
  getOperation().walk(

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

        const Type element_type =