};

// If the output is produced by a callop, will return the callop, otherwise,
// will return nullptr.
inline func::CallOp GetProducerCallOpOrNull(Value output) {
  Operation* output_op = output.getDefiningOp();
  if (output_op != nullptr && llvm::isa<func::CallOp>(output_op)) {
    return llvm::cast<func::CallOp>(output_op);
  }
  return nullptr;
}

class PickSubgraphsPass

    }
  }

  // Adds the defining op of index-th operand of op to the work list.
  void AddOperandToList(Operation* op, const int index) {
    if (Operation* operand_op = op->getOperand(index).getDefiningOp();
        operand_op != nullptr) {
      work_list_.push_back(operand_op);
    }
  }

  // Returns the quantization params for the bias input from the non-bias

        should_add_op = false;
        SmallVector<Operation*> all_descendants;
        for (Value v : operands) {
          if (defined_values.contains(v)) continue;
          if (ShouldAddOpToSubgraph(v.getDefiningOp(), reverse_subgraph,
                                    ops_to_add, all_descendants)) {
            should_add_op = true;
            break;
          }
        }
        if (should_add_op) {

llvm::SmallVector<Operation*> GetPredecessors(Operation* op) {
  llvm::SmallVector<Operation*> predecessors;
  for (auto operand : op->getOperands()) {
    if (Operation* pred = operand.getDefiningOp()) {
      pred->walk([&](mlir::Operation* opinexecutor) {
        predecessors.push_back(opinexecutor);
      });
    }
  }
  return predecessors;
}

    host_op->walk([&](Operation* op) {
      // Add any operations that provide variant inputs to the cluster.
      for (auto value : op->getOperands()) {
        Operation* input_defining_op = value.getDefiningOp();
        if (IsVariant(value) && input_defining_op &&
            !HasOutsideCompiledAncestor(input_defining_op) &&
            !input_defining_op->hasAttrOfType<StringAttr>(

    // `tf.TPUPartitionedInputV2` op is input to the i-th logical device.
    if (auto partitioned_input =
            llvm::dyn_cast_or_null<mlir::TF::TPUPartitionedInputV2Op>(
                input_value.getDefiningOp())) {
      if (UnsupportedPartitionedShardingType(input_sharding_type))
        return cluster_func->emitOpError()
               << "unsupported input sharding type "

// Returns a non-null value and can return `value` if backtracking is not
// possible.
Value BacktrackAnalysis::BacktrackValue(Value value) {
  while (Operation* op = value.getDefiningOp()) {
    int res_index = mlir::cast<OpResult>(value).getResultNumber();
    if (auto graph = dyn_cast<tf_executor::GraphOp>(op)) {
      value = graph.GetFetch().getOperand(res_index);

  for (Value fetch_operand : fetch.getOperands()) {
    if (mlir::isa<mlir::tf_executor::ControlType>(fetch_operand.getType())) {
      Operation* defining_op =
          GetIslandInnerOpOrSelf(fetch_operand.getDefiningOp());
      auto node_it = nodes_.find(defining_op);
      TF_RET_CHECK(node_it != nodes_.end());
      control_ret_nodes->insert(node_it->second);
    }
  }
  return absl::OkStatus();
}

  for (Value fetch_operand : fetch.getOperands()) {
    if (mlir::isa<mlir::tf_executor::ControlType>(fetch_operand.getType())) {
      Operation* defining_op =
          GetIslandInnerOpOrSelf(fetch_operand.getDefiningOp());
      auto node_it = nodes_.find(defining_op);
      TF_RET_CHECK(node_it != nodes_.end());
      control_ret_nodes->insert(node_it->second);
    }
  }
  return absl::OkStatus();
}

// returns an error.
LogicalResult GetConstShapeValue(Value shape_value,
                                 llvm::SmallVector<int64_t, 8>* shape) {
  auto shape_op = shape_value.getDefiningOp();
  if (!shape_op) return failure();
  auto shape_const_op = llvm::dyn_cast<TF::ConstOp>(shape_op);
  if (!shape_const_op) return failure();
  for (const auto& v : shape_const_op.getValue().getValues<APInt>()) {