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

    tf_device.return %1 : tensor<2x5x16xi32>
  }) : () -> (tensor<2x5x16xi32>)
  func.return %0: tensor<2x5x16xi32>
}


// Tests that composite tf.ResourceApplyCenteredRMSProp operation is decomposed.

// CHECK-LABEL: func @decompose_resource_apply_centered_RMS_prop

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>
}

  %0 = "chlo.broadcast_add"(%arg0, %arg1) {broadcast_dimensions = array<i64: 1, 2, 3>} : (tensor<4x1x1xi32>, tensor<4x4x4x4xi32>) -> tensor<4x4x4x4xi32>
  func.return %0 : tensor<4x4x4x4xi32>
}

// CHECK-LABEL:   func @unsupported_broadcast_add
// CHECK: chlo.broadcast_add

// -----

// CHECK-LABEL: @unsorted_segment_sum
// CHECK-SAME: [[DATA:%.*]]: tensor<8x16x64xf32>
// CHECK-SAME: [[SI:%.*]]: tensor<8x16xi32>
func.func @unsorted_segment_sum(%data: tensor<8x16x64xf32>, %segment_ids : tensor<8x16xi32>) -> (tensor<4x64xf32>) {
  %num_segments = "tf.Const"() {value = dense<4> : tensor<i32>} : () -> tensor<i32>

}

// -----

// 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