mlir::tf_device::ClusterFuncOp cluster_func, mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list) {
  // Initialize the input list for each logical devices.
  input_list->reserve(num_cores_per_replica);
  for (int i = 0; i < num_cores_per_replica; ++i)
    input_list->emplace_back(llvm::SmallVector<mlir::Value, 4>());

  llvm::SmallVector<mlir::Value, 4> cluster_func_inputs(

  // i-th logical core.
  llvm::SmallVector<llvm::SmallVector<mlir::Value, 4>, 4> input_list;
  builder->setInsertionPoint(*new_parallel_execute);
  auto result = tensorflow::ExtractInputsForLogicalDevices(
      num_cores_per_replica, cluster_func, builder, &input_list);
  if (failed(result)) return failure();

  const bool replicated = tpu_devices.size() != 1;

template <typename BatchMatMulOpType>
LogicalResult ConvertTFBatchMatMulOp<BatchMatMulOpType>::matchAndRewrite(
    BatchMatMulOpType op, PatternRewriter& rewriter) const {
  Value input_lhs = op.getX();
  Value input_rhs = op.getY();

  if (!mlir::isa<RankedTensorType>(input_lhs.getType())) {
    // LHS must be a ranked tensor type
    return failure();
  }
  if (!mlir::isa<RankedTensorType>(input_rhs.getType())) {

  if (auto input_result = mlir::dyn_cast<mlir::OpResult>(src)) {
    auto* input_inst = GetIslandInnerOpOrSelf(input_result.getOwner());
    // Replaces the input node with NextIteration sink if it is a NextIteration
    // source.
    if (auto next_iter_source =
            llvm::dyn_cast<mlir::tf_executor::NextIterationSourceOp>(
                input_inst))
      input_inst = next_iter_source.GetSink();

  if (auto input_result = mlir::dyn_cast<mlir::OpResult>(src)) {
    auto* input_inst = GetIslandInnerOpOrSelf(input_result.getOwner());
    // Replaces the input node with NextIteration sink if it is a NextIteration
    // source.
    if (auto next_iter_source =
            llvm::dyn_cast<mlir::tf_executor::NextIterationSourceOp>(
                input_inst))
      input_inst = next_iter_source.GetSink();

  };

  if (!can_convert_to_fully_connected(op)) {
    return false;
  }

  Value input_lhs = get_real_input_value(op.getX());
  Value input_rhs = get_real_input_value(op.getY());

  Value legalized_lhs =
      op.getAdjX() ? create_z_x_transpose_op(input_lhs) : input_lhs;

  // The rhs need to be transposed if adj_y == false AND this matmul will be
  // legalized to tfl.fully_connected