}
  while_op.replaceAllUsesWith(
      new_while.getResults().take_front(while_op.getNumResults()));
  while_op.erase();
  return success();
}

LogicalResult HandleIfRegionOp(
    TF::IfRegionOp if_op, ModuleOp module,
    llvm::SmallDenseMap<Value, SizeInfo>* buffer_to_size,
    llvm::StringMap<PartitionedCallDecompositionInfo>*
        decomposed_partitioned_call_callees) {
  // Rewrite the branches.

      op);
}

// Returns true if the given op is one of ops supported to have communication
// subcomputation in the TF/XLA bridge.
bool SupportsCommunicationComputation(Operation* op) {
  return isa<TF::IfRegionOp, TF::WhileRegionOp, TF::CaseRegionOp,
             TF::XlaCallModuleOp, TF::StatefulPartitionedCallOp,
             TF::PartitionedCallOp, TF::LegacyCallOp>(op);
}

#define GEN_PASS_DEF_PREPARETPUCOMPUTATIONFORTFEXPORTPASS

                       IRMapping &valueMapping) const final {
    // Allow inlining in regions attached to region based control flow
    // operations only if the src region is a single block region
    return isa<IfRegionOp, CaseRegionOp, WhileRegionOp>(dest->getParentOp()) &&
           llvm::hasSingleElement(*src);
  }

  // Returns true if its legal to inline a TF operation `op` into the `dest`
  // region.

        AddRegionSideEffectsForOp(branch_func.getBody(), op);
      }
    } else if (isa<tf_device::LaunchOp, tf_device::ClusterOp,
                   tf_executor::IslandOp, tf_executor::GraphOp, IfRegionOp,
                   CaseRegionOp>(op)) {
      for (Region& region : op->getRegions()) {
        AddRegionSideEffectsForOp(region, op);
      }
    } else if (auto xla_call_module_op = dyn_cast<XlaCallModuleOp>(op)) {

      (*arguments_to_erase)[callee].push_back(argument_index - 1);
      work_list->push_back(
          std::make_pair(&callee.getBody(), argument_index - 1));
    }
  } else if (auto if_op = dyn_cast<TF::IfRegionOp>(user_op)) {
    (*arguments_to_erase)[if_op].push_back(argument_index);
    for (auto callee : {&if_op.getThenBranch(), &if_op.getElseBranch()}) {
      work_list->push_back(std::make_pair(callee, argument_index));
    }

      if (while_op.cond_function().walk(check_while_cond).wasInterrupted())
        return WalkResult::interrupt();
      result = CanonicalizeFunctionalWhile(while_op);
    } else if (isa<TF::IfRegionOp, TF::CaseRegionOp, tf_device::ClusterOp>(
                   op)) {
      result = CanonicalizeRegionIfCaseCluster(op);
    } else if (auto while_region = dyn_cast<TF::WhileRegionOp>(op)) {

      // the bug. b/195583695 describes the motivation of this change.
      // See b/216355804 how to reproduce the bug regarding tf.RandomUniform Op

      // Conditional ops
      TypeID::get<TF::IfRegionOp>(),
      TypeID::get<TF::WhileRegionOp>(),
      TypeID::get<TF::CaseRegionOp>(),
      TypeID::get<TF::YieldOp>(),
  };
  return *ops;
}

bool IsOpTypeAllowedTf2XlaFallback(const TypeID& type_id) {

  supported_ops->insert(
      OperationName(mlir::TF::CaseRegionOp::getOperationName(), context));
  supported_ops->insert(
      OperationName(mlir::TF::IfRegionOp::getOperationName(), context));
  supported_ops->insert(
      OperationName(mlir::TF::InplaceAddOp::getOperationName(), context));
  supported_ops->insert(

                            decomposed_partitioned_call_callees))) {
        return failure();
      }
    } else if (llvm::isa<TF::WhileRegionOp>(op) ||
               llvm::isa<TF::IfRegionOp>(op) ||
               llvm::isa<TF::CaseRegionOp>(op)) {
      if (failed(
              HandleRegionControlFlowOps(op, module, data_var_to_size_var,

    } else if (auto if_op = dyn_cast<IfOp>(op)) {
      AnalyzeFunctionalCaseOrIfOp(
          if_op, {if_op.then_function(), if_op.else_function()},
          backtrack_analysis);
    } else if (llvm::isa<CaseRegionOp, IfRegionOp>(op)) {
      AnalyzeRegionCaseOrIfOp(op, backtrack_analysis);
    } else if (auto call = dyn_cast<CallOpInterface>(op)) {
      func::FuncOp func = dyn_cast_or_null<func::FuncOp>(