auto island_op = llvm::dyn_cast<IslandOp>(*iter);
    if (llvm::dyn_cast_or_null<TF::IdentityOp>(island_op.GetBody().front())) {
      if (island_op.getOutputs().use_empty()) {
        tpu_ops.erase(iter);
        return true;
      }
      std::vector<IslandOp> ops;
      if (incoming) {
        collect_output_users_islands(island_op, ops);
      } else {

    Value control_barrier = control_barrier_worklist.front();
    control_barrier_worklist.pop();

    // We can only erase control barriers whose uses have been erased as well.
    if (!control_barrier.use_empty()) continue;

    // Only values defined by IslandOp were inserted in the worklist.
    IslandOp current_island = cast<IslandOp>(control_barrier.getDefiningOp());

        return std::nullopt;
      });
  if (!elem_type) return std::nullopt;
  return llvm::to_vector<8>(elem_type->getShape());
}

void ReplaceAllUsesWithCast(Value old_val, Value new_val) {
  if (old_val.use_empty()) return;
  auto cast_op =
      OpBuilder(old_val.getDefiningOp())
          .create<tensor::CastOp>(old_val.getLoc(), old_val.getType(), new_val);
  old_val.replaceAllUsesWith(cast_op);
}

      // Now check if the operand is unused in both regions as well as the
      // result. If so, mark it for removal.
      if (body_block.getArgument(op_idx).use_empty() &&
          cond_block.getArgument(op_idx).use_empty() &&
          while_op.getResult(op_idx).use_empty()) {
        removed_operand.set(op_idx);
        new_num_operands--;
      }
    }

    if (new_num_operands == old_num_operands) return failure();

  LogicalResult matchAndRewrite(IslandOp op,
                                PatternRewriter &rewriter) const override {
    if (op.getNumOperands() != 0 || op.getNumResults() != 2 ||
        !op.getControl().use_empty() ||
        !HasSingleOpInBlock<YieldOp>(&op.GetBody()))
      return failure();

    rewriter.replaceOp(op, {op.GetYield().getOperand(0), nullptr});

    return success();
  }
};

    return;
  }

  auto value = values[0];
  auto replacement = replacements[0];
  auto tuple_type = mlir::dyn_cast<TupleType>(value.getType());
  if (!tuple_type) {
    if (!value.use_empty()) {
      auto new_element = builder.create<GetTupleElementOp>(replacement.getLoc(),
                                                           replacement, 0);
      value.replaceAllUsesWith(new_element.getResult());
    }

      op->getResult(0).replaceAllUsesWith(op->getOperand(0));

    // Leftover TPUReplicatedInput/TPUReplicatedOutput that are not of
    // `num_replicas` to 1.
    if (!op->use_empty()) {
      op->emitOpError() << "is expected to have no uses, but it is operand#"
                        << op->use_begin()->getOperandNumber() << " of "
                        << *op->use_begin()->getOwner();

  func.walk([&](quantfork::QuantizeCastOp q) {
    auto input_op = q.getArg().getDefiningOp();
    if (auto stats = dyn_cast_or_null<quantfork::StatisticsOp>(input_op)) {
      q.setOperand(stats.getArg());
      if (stats.use_empty()) stats.erase();
    }
  });

  // Step 1: forward pass: propagate any value scales which are not produces
  // by `SameOperandsAndResultsScale`. Additionally, remove the value scales

      auto value = while_op.getOperand(while_index);
      if (body_block.getArgument(arg_index).use_empty() &&
          cond_block.getArgument(arg_index).use_empty() &&
          // Note: since we are not erasing results, need to use while_index
          // to check if the corresponding result is unused.
          while_op.getResult(while_index).use_empty()) {
        unchanged = false;
        body_block.eraseArgument(arg_index);

      if (failed(mergeDequantizeOpFollowingQuantizedFunction(call_op, args,
                                                             rewriter))) {
        return failure();
      }
    }
    if (call_op->use_empty()) return success();

    DenseMap<Value, quantfork::StorageCastOp> replace_map;
    rewriter.setInsertionPointAfter(call_op);

    SmallVector<Type, 4> result_types;