// CHECK: [[UNPACK:%.*]]:3 = "tf.Unpack"(%arg0) <{axis = 0 : i64}> : (tensor<3x2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>)
// CHECK: [[SCALAR_ZERO:%.*]] = arith.constant dense<0> : tensor<i32>
// CHECK: [[CONCAT:%.*]] = "tf.Concat"([[SCALAR_ZERO]], [[UNPACK]]#0, [[UNPACK]]#1, [[UNPACK]]#2) : (tensor<i32>, tensor<2x2xf32>, tensor<2x2xf32>, tensor<2x2xf32>) -> tensor<?x2xf32>

// CHECK-LABEL: broadcast_to_to_reshape
func.func @broadcast_to_to_reshape(%arg0: tensor<4x4x4xf32>, %arg1 : tensor<4xi32>) -> tensor<1x4x4x4xf32> {
  %0 = "tfl.broadcast_to"(%arg0, %arg1) : (tensor<4x4x4xf32>, tensor<4xi32>) -> tensor<1x4x4x4xf32>
  // CHECK: "tfl.reshape"
  // CHECK-SAME: (tensor<4x4x4xf32>, tensor<4xi32>) -> tensor<1x4x4x4xf32>
  func.return %0 : tensor<1x4x4x4xf32>
}

  // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32>

//     : (tensor<2xf32>, tensor<2xi64>) -> tensor<1x2xf32>
//   %inp1 = "tf.Reshape"(%arg1, %shape)
//     : (tensor<2x2x2xf32>, tensor<2xi64>) -> tensor<4x2xf32>
//   %items0 = "tf.Unpack"(%[[INP0]]) {axis = 0 : i64}
//     : (tensor<1x2xf32>) -> tensor<2xf32>
//   %items1:4 = "tf.Unpack"(%[[INP1]]) {axis = 0 : i64}
//     : (tensor<4x2xf32>) -> (tensor<2xf32>, tensor<2xf32>, tensor<2xf32>,
//     tensor<2xf32>)

// Note: Argument quantization sequence omitted.
// CHECK: stablehlo.dot_general %{{.+}}, %{{.+}}, contracting_dims = [1] x [0] : (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>

// Note: Result dequantization sequence omitted.
// CHECK: return %{{.+}} : tensor<1x3xf32>

// -----

// Tests that a simple dot_general (lifted as a function) with CustomAggregators

        %5:2 = "tf.C"(%4) : (tensor<2xi32>) -> (tensor<2xi32>, tensor<2xi32>)
        tf_device.return %5#0, %5#1 : tensor<2xi32>, tensor<2xi32>
      }) {num_cores_per_replica = 1, topology =  "", device_assignment =  []} : () -> (tensor<2xi32>, tensor<2xi32>)
      tf_device.return %2#1 : tensor<2xi32>
    }
    func.return %1 : tensor<2xi32>
  }

    %cst_1 = "tf.Const"() {value = dense<[2, 1024]> : tensor<2xi32>} : () -> tensor<2xi32>
    // Original identity preserved.
    %cst_identity = "tf.Identity"(%cst_0) {device = ""} : (tensor<1024x2xf32>) -> tensor<1024x2xf32>
    %0 = "tf.Reshape"(%cst_identity, %cst_1) : (tensor<1024x2xf32>, tensor<2xi32>) -> tensor<2x1024xf32>

    // CHECK:      %[[ONE:.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}>
    // CHECK:      %[[RES_READ_VAL:[0-9]*]] = "tf.ReadVariableOp"
    // CHECK-SAME: (tensor<*x!tf_type.resource<tensor<2x8xi32>>>) -> tensor<2x8xi32>
    // CHECK:      "tf.AddV2"(%[[RES_READ_VAL]], %[[ONE]])
    // CHECK-SAME: (tensor<2x8xi32>, tensor<i32>) -> tensor<2x8xi32>
    // CHECK:      "tf.AssignVariableOp"

//                (tensor<4x2xf32>, tensor<4x2xf32>, tensor<4x2xf32>)
//
// will be converted into:
//
//   %0 = "mhlo.slice"(%input) {
//             limit_indices = dense<[4, 2]> : tensor<2xi64>,
//             start_indices = dense<0> : tensor<2xi64>,
//             strides = dense<1> : tensor<2xi64>} :
//        (tensor<4x6xf32>) -> tensor<4x2xf32>

}

// Returns a RankedTensorType which is similar to `input_type` but replaces the
// dimension size of `dim` with `dim_size`.  For example,
// `SubstituteRankedTensorTypeDimSize(tensor<3x4xi32>, 1, 2)` returns
// `tensor<3x2xi32>`.
static RankedTensorType SubstituteRankedTensorTypeDimSize(
    RankedTensorType input_type, int64_t dim, int64_t dim_size) {
  auto shape = input_type.getShape().vec();
  shape[dim] = dim_size;