Block& front = region->front();
    auto old_return = front.getTerminator();
    assert(old_return->getNumOperands() == op_->getNumResults());
    auto new_return_operands = llvm::to_vector<4>(old_return->getOperands());
    new_return_operands.resize(num_new_results_);

    // initialize return values for written resources to be the hoisted reads.
    for (Value resource : written_resources_) {

void RemovePassthroughOp(Block &block) {
  for (auto &op : llvm::make_early_inc_range(block)) {
    if (isa<TF::IdentityOp, TF::IdentityNOp>(op) || IsCastOfResource(op)) {
      op.replaceAllUsesWith(op.getOperands());
      op.erase();
    }
  }
}

using LocalVarOp = std::variant<TF::VarHandleOp, TF::MlirLocalVarOp>;

Value LocalVarOp_resource(LocalVarOp &op) {

    // which should be removed. So the replaced PartitionedCall op should keep
    // its original arguments except for the last element.
    SmallVector<Value> args = call_op.getOperands().drop_back();

    rewriter.replaceOpWithNewOp<TF::PartitionedCallOp>(
        call_op, float_func.getResultTypes(), args, f_attr);
    return success();
  }
};

template <typename OpT>
OpT AddOperandAndRewriteAs(Operation* op, Value operand, NamedAttrList attr,
                           OpBuilder* builder) {
  builder->setInsertionPoint(op);
  auto operands = llvm::to_vector<4>(op->getOperands());
  operands.push_back(operand);
  auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
                                     operands, attr.getAttrs());

                                   SmallVector<TransposeOp, 2>* transpose_ops) {
  for (auto it = transpose_ops->begin(); it != transpose_ops->end(); ++it) {
    auto tranpose_op = *it;
    for (auto tranpose_operand : tranpose_op.getOperands()) {
      auto ranked_tranpose_type =
          mlir::dyn_cast_or_null<RankedTensorType>(tranpose_operand.getType());
      if (!ranked_tranpose_type) continue;
      if (ranked_tranpose_type.getRank() == permutation.size() &&

      }
      // Map the block arguments to the yield results.
      for (auto iter_op : llvm::zip(for_op.getRegionIterArgs(),
                                    terminator_op->getOperands())) {
        mapping.map(std::get<0>(iter_op), std::get<1>(iter_op));
      }
      rewriter.eraseOp(terminator_op);
    }
    SmallVector<Value, 4> returned;
    for (Value arg : for_op.getRegionIterArgs()) {

      ancestor[child] = &op;
      for (Operation* predecessor : ctrlPredecessors(child)) {
        if (predecessor->getBlock() == &block) {
          incoming_ctrl_edges++;
        }
      }
      for (Value v : child->getOperands()) {
        if (v.getParentBlock() == &block) {
          incoming_data_edges++;
        }
      }
    });
    remaining_incoming_data_edges[&op] = incoming_data_edges;

        func_op.getRegion().getBlocks().front().getTerminator());
    if (!return_op) return failure();
    auto req_op = llvm::dyn_cast_or_null<mlir::stablehlo::UniformQuantizeOp>(
        return_op.getOperands()[0].getDefiningOp());
    if (!req_op) return failure();

    // Create a new func.call op with f32 output.
    auto new_call_op = call_op.clone();
    new_call_op->getResult(0).setType(

    // Build the newly fused operation to replace the batch norm
    OperationState state(batch_norm.getLoc(),
                         _FusedBatchNormExOp::getOperationName());
    state.addOperands(batch_norm.getOperands());
    if (side_input) state.operands.push_back(side_input);
    state.addTypes(batch_norm.getResultTypes());
    state.addAttributes(batch_norm->getAttrs());
    Operation *op = rewriter.create(state);

  branch_types.reserve(branches.size());

  if (llvm::any_of(op->getOperands(),
                   [](Value value) { return value == nullptr; }))
    return op->emitOpError("operation has null operand");

  // Functions have one less operand compared to op as first operand is elided
  // (`cond` of `tf.If` and `branch_index` of `tf.Case`).
  TypeRangeWithDesc input{op->getOperands().drop_front().getTypes(), "input"};