LogicalResult matchAndRewrite(quantfork::DequantizeCastOp op,
                                PatternRewriter& rewriter) const override {
    auto input_op = op.getArg().getDefiningOp();
    if (auto q = llvm::dyn_cast_or_null<quantfork::QuantizeCastOp>(input_op)) {
      if (!q->getAttr(kVolatileOpAttrName)) return failure();

      // If the quantize op is a requantize op, it is being used in other scale

    std::vector<std::vector<int>> supported_block_size;
    for (int operand : sparse_operands) {
      auto* op = sparse_op.getOperation();
      auto value = op->getOperand(operand);

      auto* inst = value.getDefiningOp();
      if (!inst) {
        continue;
      }

      // There could be a Dequantize op after the weight tensor in cases like
      // fp16 post-training quantization. We need to get the weight from the

  WalkResult ready_to_schedule = op.walk([&](Operation* nested_op) {
    for (Value operand : nested_op->getOperands()) {
      Operation* defining_op = operand.getDefiningOp();
      if (!defining_op) continue;
      Operation* producer_in_block = block->findAncestorOpInBlock(*defining_op);
      if (producer_in_block && producer_in_block != &op &&

  auto assign_var = dyn_cast<TF::AssignVariableOp>(use.getOwner());
  if (!assign_var) return nullptr;

  if (use.get() != assign_var.getValue()) return nullptr;

  auto* resource_handle_op = assign_var.getResource().getDefiningOp();
  if (resource_handle_op == parallel_execute) return nullptr;

  if (resource_handle_op &&
      resource_handle_op->getBlock() ==
          parallel_execute.getOperation()->getBlock() &&

  }

  for (int i = 0; i < original_launch_op->getNumResults(); i++) {
    Value operand = original_launch_op.GetBody().getTerminator()->getOperand(i);
    auto defining_op = operand.getDefiningOp();
    if (defining_op != nullptr &&
        defining_op->getParentOp() == reduced_launch_op.getOperation()) {
      auto iter = reduced_return_operand_map.find(operand);

  if (!return_op) return failure();
  if (return_op.getNumOperands() != 1) return failure();

  ReductionOp reduce_op = dyn_cast_or_null<ReductionOp>(
      return_op.getOperands().front().getDefiningOp());
  if (!reduce_op) return failure();
  if (reduce_op.getLhs() != body.getArgument(0) ||
      reduce_op.getRhs() != body.getArgument(1))
    return failure();

  return success();
}

  for (uint64_t i = 0, e = op->getNumOperands(); i < e; ++i) {
    // Check that the i'th operand is a broadcast.
    auto broadcast = llvm::dyn_cast_or_null<TF::BroadcastToOp>(
        op->getOpOperand(i).get().getDefiningOp());
    if (!broadcast) continue;

    // Check that the operand of the broadcast has fully defined shape.
    auto broadcast_arg_type = mlir::dyn_cast_or_null<RankedTensorType>(

  // An operation is ready to be scheduled if all its operands are ready. An
  // operand is ready if:
  const auto is_ready = [&](Value value, Operation *top) {
    Operation *parent = value.getDefiningOp();
    // - it is a block argument,
    if (parent == nullptr) return true;
    Operation *ancestor = block->findAncestorOpInBlock(*parent);
    // - it is an implicit capture,

    }
    alias_info.output_index = result.getResultNumber();
  }

  for (auto& operand : cluster_func->getOpOperands()) {
    auto read_op = llvm::dyn_cast_or_null<TF::ReadVariableOp>(
        operand.get().getDefiningOp());
    if (!read_op) continue;
    if (!read_op->hasOneUse()) continue;
    auto it = resource_alias_info_map.find(read_op.getResource());
    if (it == resource_alias_info_map.end()) continue;

        for (Value v : user->getOperands()) {
          if (tensorflow::TypeValidForXLA(v.getType()) &&
              v.getDefiningOp() == op &&
              !llvm::isa<mlir::tf_device::ReturnOp>(user))
            external_outputs.insert(v);
          if (v.getDefiningOp() == op &&
              llvm::isa<mlir::tf_device::ReturnOp>(user))
            tmp_host_outputs.push_back(v);
        }
      }
    }