for (Operation& head_outside_compiled_op : *launch_block) {
    for (Value result : head_outside_compiled_op.getResults()) {
      bool has_external_uses = false;
      for (Operation* user : result.getUsers()) {
        if (OpInBlock(user, launch_block)) continue;
        has_external_uses = true;
        break;
      }
      if (has_external_uses) {
        launch_results.push_back(result);

    if (IsConstantOrNone(op)) {
      continue;
    }
    for (Value output : op->getResults()) {
      bool output_consumed_outside_subgraph = false;
      for (Operation* consumer : output.getUsers()) {
        if (llvm::all_of(partition_ops, [&](Operation* op) {
              return !op->isAncestor(consumer);
            })) {
          output_consumed_outside_subgraph = true;
        }
      }

                                PatternRewriter &rewriter) const override {
    // Only rewrite if the Broadcast has only one consumer.
    if (!op.getOutput().hasOneUse()) return failure();

    Operation *user = *op.getOutput().getUsers().begin();

    auto reshape_op = llvm::dyn_cast_or_null<ReshapeOp>(user);
    if (!reshape_op) return failure();

    auto reshape_type =
        mlir::cast<ShapedType>(reshape_op.getOutput().getType());

bool CanMove(Operation* op) {
  auto launch_op = llvm::dyn_cast<tf_device::LaunchOp>(op->getParentOp());
  if (!launch_op) return false;
  for (Value result : op->getResults()) {
    for (Operation* user : result.getUsers()) {
      if (user != launch_op.GetBody().getTerminator()) return false;
    }
  }
  return true;
}

// Get the new launch op results. This is the results if the copy op is removed
// from the old launch op.

        llvm::dyn_cast_or_null<quantfork::StatisticsOp>(
            op.getOperand(input_index).getDefiningOp()),
        llvm::dyn_cast_or_null<quantfork::StatisticsOp>(
            *op.getResult().getUsers().begin()),
    };

    if (!stats_ops[0] || !stats_ops[1]) {
      return failure();  // Already converted to Q-DQ pair.
    }

    llvm::SmallVector<llvm::APFloat, 4> min_max_values;

    Type cast_type = quant_type.castFromExpressedType(expressed_type);

    // Insert DQ-op if it does not exist yet. Otherwise, just rewire without
    // creating a new DQ-op.
    for (auto connected_op : op->getUsers()) {
      auto q_op =
          llvm::dyn_cast_or_null<quantfork::QuantizeCastOp>(connected_op);
      if (q_op && q_op.getType() == cast_type) {
        auto dq_op = llvm::cast<quantfork::DequantizeCastOp>(

    return HasValidWeightOnlyPtqMethod(method.weight_only_ptq(), rank);
  }

  void rewrite(Operation* op, PatternRewriter& rewriter) const override {
    Operation* quantizable_op = *op->getUsers().begin();
    DenseFPElementsAttr attr;
    matchPattern(op->getResult(0), m_Constant(&attr));

    Method method = GetQuantizationMethodOrDefault(quantizable_op);

    return;
  }
  if (!tensor_type.getElementType().isF32()) return;

  // If the result is consumed by a quantize op, it has been quantized.
  if (value.hasOneUse() &&
      llvm::isa<TFL::QuantizeOp>(*value.getUsers().begin()))
    return;

  // Add this result to the list to apply the default value.
  values->push_back(value);
}

void DefaultQuantParamsPass::QuantizeValue(OpBuilder builder, Value value,

    if (!convert_op.getResult().hasOneUse()) {
      return false;
    }
    auto other_convert_op = dyn_cast_or_null<OtherConvertOp>(
        *convert_op.getResult().getUsers().begin());
    return other_convert_op &&
           converter.isLegal(other_convert_op.getResult().getType());
  }
  return true;
}

class BFloat16TypeConversionTarget : public ConversionTarget {
 public:

  // the operands.
  for (Value result : parallel_execute.getExecuteOutputs()) {
    if (!result.hasOneUse()) continue;
    auto assign_var =
        llvm::dyn_cast<TF::AssignVariableOp>(*result.getUsers().begin());
    if (!assign_var || assign_var.getValue() != result) continue;
    auto partitioned_input =
        llvm::dyn_cast_or_null<TF::TPUPartitionedInputV2Op>(
            assign_var.getResource().getDefiningOp());