llvm::SmallVector<Operation*, 4> candidates;
    for (Value value : values) {
      if (incoming) {
        candidates = {value.getDefiningOp()};
      } else {
        candidates.assign(value.getUsers().begin(), value.getUsers().end());
      }
      for (Operation* candidate_op : candidates) {
        if (cluster_ops.contains(candidate_op) ||
            cluster_dependent_ops.contains(candidate_op)) {

    if (!type || !type.getElementType().isF32()) {
      return failure();
    }
    return success(
        op->hasOneUse() &&
        IsWeightQuantizableFunction(*op->getUses().begin(), type.getRank()));
  }

  // Checks if the operand is second operand of `tf.XlaCallModule` op for
  // `stablehlo.convolution` or `stablehlo.dot_general` with fully_quantizable
  // trait.

    if (!type || !type.getElementType().isF32()) return false;

    Value value = op.getResult();

    // Check whether dynamic range quantization can be applied.
    for (auto& use : value.getUses()) {
      Operation* user = use.getOwner();
      int operand_num = use.getOperandNumber();
      std::unique_ptr<OpQuantSpec> spec = GetTFOpQuantSpec(user);

    // Unsure if this check is required.
    if (state.pos != pos) {
      return;
    }
  }
  if (pos == RequantizeState::ON_OUTPUT) {
    Operation* user = value.getUses().begin().getUser();
    if (isa<quantfork::QuantizeCastOp>(user)) {
      // The requantize op is inserted between `quantize` and `dequantize` ops.
      value = user->getResult(0);
      builder_.setInsertionPointAfter(user);

}

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();
      assert(body);
      auto type_from_body = GetElementTypeFromAccess(

Value GetResultInBlock(Value value, llvm::DenseSet<Block*> target_blocks) {
  while (target_blocks.find(value.getParentBlock()) == target_blocks.end()) {
    Value new_value = nullptr;
    for (OpOperand& operand : value.getUses()) {
      Operation* owner = operand.getOwner();
      if (!llvm::isa<tf_device::ReturnOp>(owner)) continue;
      Operation* parent = owner->getParentOp();  // op that owns the "return"

      }
      return_operands.push_back(resource.live_value);
      result_types.push_back(resource.live_value.getType());
    }

    if (is_var_handle && !resource.read) {
      assert(block.getArgument(index).getUses().empty());
      argument_indices_to_remove.push_back(index);
    } else {
      if (is_var_handle) {
        // Add resource_name attribute to VarHandleOp read.
        function.setArgAttr(

    for (Operation* user : island.getControl().getUsers()) {
      DCHECK_EQ(user->getParentOp(), graph);
      try_update_current_candidate(user);
    }

    // Check island data results.
    Block& graph_body = llvm::cast<GraphOp>(graph).GetBody();
    for (Value result : island.getOutputs()) {
      for (Operation* user : result.getUsers()) {
        Operation* def = graph_body.findAncestorOpInBlock(*user);

  // GraphOp, YieldOp and NextIterationSourceOp don't have control inputs so
  // exclude them below.
  for (Value out : island_op.getOutputs()) {
    for (auto& use : out.getUses()) {
      Operation* owner = use.getOwner();
      if (owner->getDialect() == island_op->getDialect() &&
          !llvm::isa<tf_executor::GraphOp, tf_executor::YieldOp,

    if (constant_op.getValue().getNumElements() > 32) continue;
    while (!constant_op.getResult().hasOneUse()) {
      auto new_constant_op = builder.clone(*constant_op.getOperation());
      constant_op.getResult().getUses().begin()->assign(
          dyn_cast<mlir::stablehlo::ConstantOp>(new_constant_op));
    }
  }
}

void ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass::
    runOnOperation() {