func.func @broadcast_multi_dim_add(%arg0: tensor<4x1x1xi32>, %arg1: tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32> {
  // CHECK-NEXT: %[[LHS_BCAST:.+]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[1, 2, 3]> : tensor<3xi64>}>
  // CHECK-NEXT: mhlo.add %[[LHS_BCAST]], %arg1
  %0 = "tf.AddV2"(%arg0, %arg1) : (tensor<4x1x1xi32>, tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32>
  func.return %0: tensor<4x4x4x4xi32>
}

// CHECK:           %[[VAL_5:.*]] = "tf.Const"() <{value = dense<[1, 0]> : tensor<2xi64>}> : () -> tensor<2xi64>
// CHECK:           %[[VAL_6:.*]] = "tf.Transpose"(%[[VAL_4]], %[[VAL_5]]) : (tensor<1x1xi32>, tensor<2xi64>) -> tensor<1x1xi32>
// CHECK:           return %[[VAL_6]] : tensor<1x1xi32>
// CHECK:         }

    %10 = stablehlo.convert %8 : (tensor<8x16x16xi8>) -> tensor<8x16x16xi32>
    %11 = stablehlo.convert %2 : (tensor<1x1x1xi8>) -> tensor<1x1x1xi32>
    %12 = stablehlo.broadcast_in_dim %11, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x16xi32>
    %13 = stablehlo.subtract %10, %12 : tensor<8x16x16xi32>  // q1 - z1
    %14 = stablehlo.convert %9 : (tensor<8x16x4xi8>) -> tensor<8x16x4xi32>

func.func @reshape(%arg0 : tensor<1x128xi32>) -> tensor<4x32x1xi32>{
  %0 = "vhlo.reshape_v1"(%arg0) : (tensor<1x128xi32>) -> tensor<4x32x1xi32>
  func.return %0 : tensor<4x32x1xi32>
}

//CHECK:func.func private @reshape(%arg0: tensor<1x128xi32>) -> tensor<4x32x1xi32> {
//CHECK-NEXT: %0 = "vhlo.reshape_v1"(%arg0) : (tensor<1x128xi32>) -> tensor<4x32x1xi32>
//CHECK-NEXT: return %0 : tensor<4x32x1xi32>
//CHECK-NEXT:}

    %1 = "tf.MatMul"(%arg0, %0) {device = "", transpose_a = false, transpose_b = false} : (tensor<1x3xi32>, tensor<3x1xi32>) -> tensor<1x1xi32>
    func.return %1 : tensor<1x1xi32>
  }
  func.func @predict(
    ) -> (tensor<0x!tf_type.string> {tf_saved_model.index_path = ["r"]})
      attributes {

func.func @eliminateExtraSelectLhs(%arg0: tensor<4x2x1xf32>, %arg1: tensor<4x2x1xi1>) -> (tensor<4x2x1xf32>) {
  %cst0 = arith.constant dense<1.0> : tensor<2x2xf32>
  %cst1 = arith.constant dense<2.0> : tensor<2xf32>
  %cst_zeros = arith.constant dense<0.0> : tensor<4x2x1xf32>

  %0 = "tfl.select_v2"(%arg1, %cst_zeros, %arg0) : (tensor<4x2x1xi1>, tensor<4x2x1xf32>, tensor<4x2x1xf32>) -> tensor<4x2x1xf32>

  func.func @binary_op_broadcast(%arg0: tensor<4x1xf32>, %arg1: tensor<4x1x4xf32>) -> tensor<4x4x4xf32> {
    // CHECK: %[[BROADCAST0:.*]] = "mhlo.broadcast_in_dim"(%arg0) <{broadcast_dimensions = dense<[1, 2]> : tensor<2xi64>}> : (tensor<4x1xf32>) -> tensor<4x4x1xf32>
    // CHECK: %[[RESHAPE0:.*]] = mhlo.reshape %[[BROADCAST0]] : (tensor<4x4x1xf32>) -> tensor<4x4xf32>

  std::string buf;
  llvm::raw_string_ostream os(buf);
  status_or.value().print(os);
  return os.str();
}

TEST(MlirConvertType, ConvertToMlirTensorType) {
  // Simple case of static shapes.
  EXPECT_EQ("tensor<4x8x16xi32>",
            ConvertToMlirString({4, 8, 16}, /*unknown_rank=*/false,
                                DataType::DT_INT32));

  // Partially known shapes.
  EXPECT_EQ("tensor<?x27x?xbf16>",

func.func @Batchmatmul2FullyconnectedBatchedY(%arg0: tensor<4x128x2xf32>) -> (tensor<4x128x1xf32>) {
  %0 = arith.constant dense<42.> : tensor<4x2x1xf32>
  %1 = "tfl.batch_matmul"(%arg0, %0) {adj_x = false, adj_y = false, asymmetric_quantize_inputs = false} : (tensor<4x128x2xf32>, tensor<4x2x1xf32>) -> tensor<4x128x1xf32>
  func.return %1 : tensor<4x128x1xf32>

  // CHECK: %[[BMM_RES:.*]] = "tfl.batch_matmul"
  // CHECK-NEXT: return %[[BMM_RES]]

}

// -----

// CHECK-LABEL: func @selectV2BroadcastAll
func.func @selectV2BroadcastAll(%arg0: tensor<8x1x1xi1>, %arg1: tensor<1x8x1xi32>, %arg2: tensor<1x1x8xi32>) -> tensor<8x8x8xi32> {
  %0 = "tf.SelectV2"(%arg0, %arg1, %arg2) : (tensor<8x1x1xi1>, tensor<1x8x1xi32>, tensor<1x1x8xi32>) -> tensor<8x8x8xi32>
  func.return %0: tensor<8x8x8xi32>
}

// -----

// CHECK-LABEL: func @selectV2DynamicRanked