next_output_index, dtype);

  std::vector<NodeBuilder::NodeOut> input_list({enter_input, next_input});
  const string unique_name = scope.GetUniqueNameForOp("Merge");
  NodeBuilder builder = NodeBuilder(unique_name, "Merge").Input(input_list);
  scope.UpdateBuilder(&builder);

  Node* merge_node;
  TF_RETURN_IF_ERROR(builder.Finalize(scope.graph(), &merge_node));

mlir::LogicalResult ExtractInputsForLogicalDevices(
    int num_cores_per_replica, mlir::tf_device::ClusterFuncOp cluster_func,
    mlir::OpBuilder* builder,
    llvm::SmallVectorImpl<llvm::SmallVector<mlir::Value, 4>>* input_list);

// Extracts a list of OpSharding that represent output sharding configuration of
// `tf_device.cluster`.
mlir::LogicalResult ParseAndValidateOutputSharding(

  LogicalResult matchAndRewrite(BatchMatMulOpType op,
                                PatternRewriter& rewriter) const override {
    Value input_lhs = op.getX();
    Value input_rhs = op.getY();

    // LHS and RHS must be a ranked tensor type
    auto lhs_type = mlir::dyn_cast<RankedTensorType>(input_lhs.getType());
    auto rhs_type = mlir::dyn_cast<RankedTensorType>(input_rhs.getType());

    if (!lhs_type || !rhs_type) return failure();

          input, permutation_tensor_op.getResult());
    };

    Value input_lhs = bmm_op.getX();
    Value input_rhs = bmm_op.getY();

    Value output_lhs =
        bmm_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