kMatMulModuleStr,
          ConfigProto::Experimental::MLIR_BRIDGE_ROLLOUT_UNSPECIFIED));
}

struct MatMulTestCase {
  std::string mat_mul_method;
};

using BatchMatMulTest = TestWithParam<MatMulTestCase>;

TEST_P(BatchMatMulTest, BatchMatMul) {
  const MatMulTestCase& test_case = GetParam();
  static constexpr char kMatMulModuleStr[] = R"(