// We need to make sure both splits & split dim are constants.
  auto splits = splitv_op.getSizeSplits().getDefiningOp();
  mlir::DenseIntElementsAttr splits_attr;
  if (!splits || !matchPattern(splits, m_Constant(&splits_attr)))
    return failure();

  auto split_dim = splitv_op.getSplitDim().getDefiningOp();
  mlir::ElementsAttr split_dim_attr;
  if (!split_dim || !matchPattern(split_dim, m_Constant(&split_dim_attr)))

  // Broadcast bias value if unmatched with output shape.
  auto bcast_op = TryCast<BroadcastInDimOp>(bias_op.getDefiningOp(),
                                            /*name=*/"broadcast_in_dim_op");

  if (failed(bcast_op)) {
    bcast_op = TryCast<DynamicBroadcastInDimOp>(
        bias_op.getDefiningOp(),
        /*name=*/"dynamic_broadcast_in_dim_op");
  }
  // Update the bias type for both static and dynamic broadcasts.

mlir::Operation* GetXlaShardingFromOperand(Value value) {
  Value value_to_visit = value;
  if (auto read_var = value_to_visit.getDefiningOp<mlir::TF::ReadVariableOp>())
    value_to_visit = read_var.getResource();

  if (auto partitioned_input =
          value_to_visit.getDefiningOp<mlir::TF::TPUPartitionedInputV2Op>()) {
    return NullUnlessSharded(partitioned_input);
  }

  return nullptr;
}

    // Only values defined by IslandOp were inserted in the worklist.
    IslandOp current_island = cast<IslandOp>(control_barrier.getDefiningOp());

    for (auto control_input : current_island.getControlInputs()) {
      ops_connected_to_fetch->insert(control_input.getDefiningOp());
      if (IsNoOpControlBarrier(control_input))
        control_barrier_worklist.push(control_input);
    }
    current_island.erase();

      if (matchPattern(quantize_operand, m_Constant(&attr))) {
        // Const-> QuantizeOp pattern will be handled separately.
        return failure();
      }
      if (Operation* quantizing_op = quantize_operand.getDefiningOp()) {
        quantizing_ops.push_back(quantizing_op);
      }
    }

    tensorflow::DataType inference_type =
        quant_params_.quant_spec.inference_type;
    bool weight_only_quantization =

                                    Value max_size, Operation* op,
                                    Type element_dtype, OpBuilder builder,
                                    Value* buffer) {
  auto max_count_op = max_size.getDefiningOp();
  if (!max_count_op) return op->emitOpError("unknown max element count");
  auto max_count_const_op = llvm::dyn_cast<TF::ConstOp>(max_count_op);
  if (!max_count_const_op) return op->emitOpError("unknown max element count");

        all_reduce.getGroupSize()
            .getDefiningOp<TF::CollectiveAssignGroupV2Op>();

    if (assign_group) {
      // Found a group assignment. Use replica_groups to represent group
      // assignment.

      if (assign_group != all_reduce.getGroupKey()
                              .getDefiningOp<TF::CollectiveAssignGroupV2Op>()) {
        return all_reduce->emitOpError()

    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 &&

  }

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