return %1#0, %1#1 : tensor<1x2xf32>, tensor<1x2xf32>
}
func.func @_func(%arg0: tensor<2x4xf32>, %arg1: tensor<4x2xf32>) -> tensor<2x2xf32> {
  %0 = "tf.MatMul"(%arg0, %arg1) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01"} : (tensor<2x4xf32>, tensor<4x2xf32>) -> tensor<2x2xf32>
  return %0 : tensor<2x2xf32>
}

// -----

  auto matmul = rewriter.create<TF::BatchMatMulV3Op>(
      loc, RankedTensorType::get(matmul_shape, result_type.getElementType()),
      lhs_flattend, rhs_flattend);

  if (result_type.hasStaticShape()) {
    auto reshaped =
        rewriter.create<mhlo::ReshapeOp>(loc, result_type, matmul.getResult());
    return reshaped.getResult();
  }

//   %1 = mhlo.reshape %param : (1xCxZ) -> CxZ
//   mhlo.dot_general %input, %1 {batch_dims = []}
// To:
//   mhlo.dot_general %input, %param {batch_dims = [0]}
//
// This usage will mostly come from tf-unroll-batch-matmul, so it's fine to only
// handle the case where batching dim is the leftmost dim.
LogicalResult ConvertReshapeDotRhsToBatchedDot(mhlo::DotGeneralOp dot,
                                               PatternRewriter &rewriter) {

// Guarantee that all FuncOp's have a single use.
std::unique_ptr<OperationPass<ModuleOp>> CreateGuaranteeAllFuncsOneUsePass();

// Optional pass which will unroll BatchMatMul and use only MatMul
std::unique_ptr<OperationPass<func::FuncOp>> CreateUnrollBatchMatMulPassPass();

// Optional pass which will map TF BatchMatMul to TF Einsum
std::unique_ptr<OperationPass<func::FuncOp>> CreateBatchMatMulToEinsumPass();

    //
    // 2. Old fashion Tensor with raw device memory pointer. This case occurs
    // when the producer is a non-XLA TF GPU kernel or function (e.g.
    // tf.matmul).
    //
    // 3. AsyncValueTensor, containing a PjRtBuffer. This is the legacy mode
    // and certain device type (e.g. TPU) still uses this path.
    AsyncValueTensor* av_tensor = AsyncValueTensor::FromTensor(tensor);

    quantization_axis = -1 : i64, quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<9x10x!tf_type.qint8>, tensor<f32>, tensor<i32>) -> tensor<9x10xf32>
  %0 = "tf.MatMul"(%input, %filter_new) {
  } : (tensor<8x9xf32>, tensor<9x10xf32>) -> tensor<8x10xf32>
  return %0 : tensor<8x10xf32>
})mlir";
  constexpr absl::string_view kProgram = R"mlir(

    ```

    The pass also works across control flow and functional calls.
  }];
}



def UnrollBatchMatMulPass : Pass<"tf-unroll-batch-matmul", "mlir::func::FuncOp"> {
  let summary = "Unroll TF BatchMatMul op into Reshape, Slice, MatMul, Pack ops.";
  let constructor = "TF::CreateUnrollBatchMatMulPassPass()";
}

def ClusterFormationPass : Pass<"tf-device-cluster-formation", "mlir::ModuleOp"> {

    if (fc_op.getFusedActivationFunction() != "NONE") return failure();

    // Only fuse multiplier if all dimensions other than the depth dimension
    // are equal to 1 since otherwise
    // `matmul(x, filter) * cst != matmul(x, filter * cst)`
    // even if `filter` and `cst` are be broadcastable.
    auto shape = cst.getType().getShape();
    if (!IsDimensionsDegenerateExceptLastOne(shape)) return failure();

  // CHECK: "mhlo.dot"(%[[UPDATED_A]], %[[UPDATED_B]])
  %0 = "tf.MatMul"(%a, %b) {transpose_a = true, transpose_b = true} : (tensor<7x5xf32>, tensor<11x7xf32>) -> tensor<5x11xf32>

  func.return %0 : tensor<5x11xf32>
}

// Verify that MatMul with ranked inputs are lowered to HLO.
// CHECK-LABEL: matmul_ranked

  func.return %0: tensor<2x3x7xf32>
}

// CHECK-LABEL: testBatchMatMulToMatMul
func.func @testBatchMatMulToMatMul(%arg0: tensor<2x3xf32>, %arg1: tensor<3x2xf32>) -> tensor<2x2xf32> {
  // CHECK: %0 = "tf.MatMul"(%arg0, %arg1) <{grad_a = false, grad_b = false, transpose_a = false, transpose_b = false}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>, tensor<3x2xf32>) -> tensor<2x2xf32>