Operation* op, llvm::function_ref<bool(Operation*)> filter) const;

  // pass the filter (returning true) will be included.
  const llvm::SmallVector<Operation*, 4>& DirectControlSuccessors(
      Operation* op) const;
  llvm::SmallVector<Operation*, 4> DirectControlSuccessors(
      Operation* op, llvm::function_ref<bool(Operation*)> filter) const;

// Returns a failure if a cycle prevents the merge from happening correctly
// without breaking dominance. The IR is left in invalid state in case of
// failure.
void CollectCandidateIslands(
    llvm::function_ref<bool(llvm::StringRef, Operation*)>
        is_op_calling_func_for_cluster,
    Operation* op, StringRef cluster_name,
    SmallPtrSet<Operation*, 16>& islands_set,
    SmallPtrSet<Operation*, 16>& wrapped_ops) {

}

// Run the TF XLA Bridge based on the input pipeline, which can be either TPU
// bridge pipeline or non TPU bridge pipeline.
tensorflow::Status RunTFXLABridge(
    ModuleOp module,
    llvm::function_ref<void(OpPassManager &pm)> pipeline_builder,
    llvm::StringRef module_name = llvm::StringRef(),
    llvm::StringRef dump_prefix = "tf_xla_bridge_v1") {
  // Explicitly check that the TensorFlow dialect can constant fold ops.

  // Returns a list of all of the mergable islands found in the graph.
  iterator_range<
      llvm::filter_iterator<SmallVector<MergedIsland>::const_iterator,
                            function_ref<bool(const MergedIsland&)>>>
  GetMergableIslands() const {
    function_ref<bool(const MergedIsland&)> filter_fn =
        [](const MergedIsland& merged_island) {
          return merged_island.islands.size() > 1;
        };

// size variables before finally changing the function type.
void ModifyFunctionSignature(
    func::FuncOp func, llvm::SmallDenseMap<Value, Value>* stack_var_to_size_var,
    llvm::function_ref<std::optional<Type>(int64_t)> arg_to_stack_type,
    llvm::function_ref<void(ArrayRef<BlockArgument>)> handle_new_size_vars =
        nullptr) {
  auto new_input_types = llvm::to_vector<8>(func.getFunctionType().getInputs());

  // Inserts constraints for multiple values.
  void Insert(ValueRange value, ValueConstraint constraint);

  // Walk all the constraints owned by this set.
  void Walk(llvm::function_ref<void(Value, ValueConstraint)> walk) const;

  // Returns the constraint of the value if it exists, or None otherwise.
  std::optional<ValueConstraint> GetConstraint(Value value) const;
  bool HasConstraint(Value value) const;

  *buffer = broadcast.getOutput();
  return success();
}

std::optional<RankedTensorType> GetElementTypeFromAccess(
    Value collection, ModuleOp module,
    llvm::function_ref<std::optional<Type>(Operation*)> infer_from_op) {
  for (auto& use : collection.getUses()) {
    if (auto while_op = llvm::dyn_cast<TF::WhileOp>(use.getOwner())) {
      auto body = while_op.body_function();

  if (it == sorted_control_predecessors_.end()) return empty_operation_set_;
  return it->second;
}

llvm::SmallVector<Operation*, 4>
SideEffectAnalysisInfo::DirectControlPredecessors(
    Operation* op, llvm::function_ref<bool(Operation*)> filter) const {
  llvm::SmallVector<Operation*, 4> result;
  auto it = sorted_control_predecessors_.find(op);
  if (it == sorted_control_predecessors_.end()) return result;

// and sizes.
void ModifyFunctionSignature(
    func::FuncOp func, Type size_type,
    llvm::SmallDenseMap<Value, SizeInfo>* buffer_to_size,
    llvm::function_ref<std::optional<Type>(int64_t)> arg_to_buffer_type,
    llvm::function_ref<bool(int64_t)> arg_buffer_size_is_fixed) {
  auto new_input_types = llvm::to_vector<8>(func.getFunctionType().getInputs());
  int64_t original_arg_count = new_input_types.size();

// adding required gradient arguments.
void ChangeFunctionInputSignature(
    func::FuncOp func,
    const llvm::SmallDenseMap<int64_t, llvm::SmallVector<string, 4>>& grads,
    llvm::function_ref<Type(int64_t)> ta_arg_buffer_type,
    llvm::function_ref<bool(int64_t)> ta_accumulate_on_write,
    llvm::SmallDenseMap<Value, TensorArrayStats>* stats) {
  int64_t original_args = func.getNumArguments();