// CHECK:  "tf.VarHandleOp"
// CHECK-NOT:  [[VARIABLE:%.*]] = "tf.ReadVariableOp"
// CHECK-NEXT:  [[KEY:%.*]], [[FUTURE:%.*]] = "tf.IfrtLoadVariable"
// CHECK-SAME:    used_by_host = true
// CHECK-NEXT:  [[MATRES:%.*]] = "tf.MatMul"(%arg0, [[FUTURE]])
// CHECK-NEXT:   [[RES:%.*]] = "tf.IfrtCall"(%arg0, [[KEY]]) <{program_id = 6515870160938153680 : i64, variable_arg_indices = [1 : i32]}>
// CHECK-NEXT:    return [[RES]], [[MATRES]] : tensor<1x1xf32>, tensor<1x1xf32>

  // May have been filtered so check for lack of failure instead of success.
  EXPECT_EQ(compilation_status.Delta(kMlirWithFallbackModeFailure), 0);
}

TEST(LegalizeTFTest, MatMul) {
  static constexpr char kMatMulModuleStr[] = R"(
  module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
    func.func @main() -> (tensor<5x11xf32>) {

}

func.func @matmul(%arg0: tensor<40x37xf32>, %arg1: tensor<37x40xf32>) -> tensor<40x40xf32> {
  %0 = "tf.MatMul"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", device = "/device:CPU:0", name = "MatMul", transpose_a = false, transpose_b = false} :
(tensor<40x37xf32>, tensor<37x40xf32>) -> tensor<40x40xf32>
  func.return %0 : tensor<40x40xf32>
// CHECK-LABEL: matmul

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

    ```

    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"> {

  // - rhs: [RHSBATCHDIMS..., RHSROWS, RHSCOLS]
  // - result: [broadcast(LHSBATCHDIMS, RHSBATCHDIMS)..., LHSROWS, RHSCOLS]
  // To perform the matmul, we need to first broadcast lhs and rhs to a common
  // set of leading dimensions before doing the actual matmul.
  // That's what the code below does.
  // In particular, we populate out_lhs and out_rhs to have dimension structure:

* `tf.config.experimental.enable_tensor_float_32_execution`

}

def TF__FusedMatMulOp : TF_Op<"_FusedMatMul", [Pure, TF_SameOperandsAndResultElementTypeResolveRef]> {
  let summary = [{
Performs a MatMul followed by a specified series of operations.
  }];

  let description = [{
The inputs to the MatMul are specified by `a` and `b`. The series of operations
that follows is specified by the `fused_ops` attribute, which is a list of TF op