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();

/// `intCalculate` is chosen to conduct the calculate.
Attribute ConstFoldBinaryOp(
    Type result_type, ArrayRef<Attribute> operands,
    llvm::function_ref<APFloat(APFloat, APFloat)> float_calculate,
    llvm::function_ref<APInt(APInt, APInt)> int_calculate) {
  // Note: All types are wrapped in tensor types in TFlite. E.g., f32 is
  // represented as tensor<f32>. So we are only handling tensor types here.

  // Update existing constraint with a new one.
  auto merged = Merge(constraint, persisted);
  return {constraints_[value] = merged, merged != persisted};
}

void ValuesConstraintSet::Walk(
    llvm::function_ref<void(Value, ValueConstraint)> walk) const {
  for (auto &kv : constraints_) walk(kv.getFirst(), kv.getSecond());
}

std::optional<ValueConstraint> ValuesConstraintSet::GetConstraint(
    Value value) const {

  for (; it != block.rend(); ++it) {
    CastOp cast = dyn_cast<CastOp>(*it);
    if (!cast || cast.getTruncate()) return std::nullopt;
  }

  return call;
}

using ArgMatcherFn = function_ref<bool(Value, Region&, Value, Region&)>;

// Returns whether the arguments of the given 2 calls are match (after looking
// through cast ops). `matcher` is the predicate used to check if two arguments
// match.

// performed in order to limit the number of values implicitly captured in this
// region before outlining.
std::unique_ptr<OperationPass<func::FuncOp>> CreateClusterConstantSinkingPass(
    llvm::function_ref<bool(tf_device::ClusterOp, ElementsAttr)> filter = {});

// Creates a pass that outlines regions of tf_device.cluster operations.
std::unique_ptr<OperationPass<ModuleOp>> CreateClusterOutliningPass();

 public:
  // Factory for creating a FlatTensorFunction.
  //
  // Params:
  //  function_def - The function_def associated with the created
  //                 FlatTensorFunction. FlatTensorFunction will register this
  //                 function_def with `ctx` on creation, and de-register it on
  //                 destruction. function_def must be non-null, but
  //                 otherwise has no lifetime requirements.

// that handles the updated value for an resource argument.
LogicalResult LiftArgRetResourcesForFunction(
    func::FuncOp func_op,
    const llvm::SmallDenseMap<int64_t, Type>& resource_data_types,
    llvm::function_ref<void(int64_t, Value)> handle_updated_arg_value) {
  RegionResourceHoister hoister(func_op);
  if (failed(hoister.Analyze())) return failure();

  // Each of these resources could be read or written in the function. If its

  return success();
}

static LogicalResult VerifyCaseOrIfOpBranchFunctions(
    SymbolTableCollection& symbol_table, Operation* op,
    ArrayRef<Attribute> branches,
    llvm::function_ref<std::string(unsigned branch_index)> branch_name) {
  SmallVector<FunctionType, 2> branch_types;
  branch_types.reserve(branches.size());

  if (llvm::any_of(op->getOperands(),