module attributes {tf.versions = {producer = 888 : i32}} {
func.func @testReshapeShapeInference(%arg0: tensor<3x4xi32>) -> tensor<*xi32> {
  %cst = arith.constant dense<[1, 6, 2]> : tensor<3xi32>
  // CHECK: "tfl.reshape"(%arg0, %cst) : (tensor<3x4xi32>, tensor<3xi32>) -> tensor<1x6x2xi32>
  %0 = "tfl.reshape"(%arg0, %cst) : (tensor<3x4xi32>, tensor<3xi32>) -> tensor<*xi32>
  func.return %0 : tensor<*xi32>
}
}

// -----

    // CHECK: mhlo.reduce
    // CHECK: mhlo.dot_general
    // CHECK: mhlo.transpose
    %0 = "tf.BatchMatMulV2"(%arg0, %arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<3x2x4xf32>) -> tensor<3x4x4xf32>
    func.return %0 : tensor<3x4x4xf32>
  }

  // CHECK-LABEL: approx_topk

  func.func @fill_with_shape_op(%arg0: tensor<3x2xi32>, %arg1: tensor<i32>) -> (tensor<*xi32>) {
    // CHECK: %[[SHAPE:.*]] = "tf.Shape"(%{{.*}}) : (tensor<3x2xi32>) -> tensor<2xi32>
    // CHECK: %[[FILL:.*]] = "tf.Fill"(%[[SHAPE]], %{{.*}}) : (tensor<2xi32>, tensor<i32>) -> tensor<3x2xi32>
    // CHECK: return %[[FILL]] : tensor<3x2xi32>
    %0 = "tf.Shape"(%arg0) : (tensor<3x2xi32>) -> tensor<2xi32>

    value {
      type: DT_BOOL
    }
  }
}

# CHECK-DAG: tf.IteratorGetNext{{.+}}-> (tensor<1x8xbf16>, tensor<2x?x16xf32>, tensor<*xf64>)
# CHECK-DAG: tf.IteratorGetNextSync{{.+}}-> (tensor<*xi16>, tensor<3x24xi32>, tensor<?x4x32xi64>)
# CHECK-DAG: tf.MultiDeviceIteratorGetNextFromShard{{.+}}-> (tensor<5x40xf16>, tensor<*xcomplex<f32>>, tensor<6x48x?xcomplex<f64>>)

    %11 = "tf.Sub"(%10, %cst_2) : (tensor<2x2x2xi32>, tensor<i32>) -> tensor<2x2x2xi32>
    %12 = "tf.Identity"(%9) : (tensor<2x2x2xi8>) -> tensor<2x2x2xi8>
    %13 = "tf.Cast"(%12) {Truncate = false} : (tensor<2x2x2xi8>) -> tensor<2x2x2xi32>
    %14 = "tf.Sub"(%13, %cst_4) : (tensor<2x2x2xi32>, tensor<i32>) -> tensor<2x2x2xi32>

      %%arg1 = "tf.Const"() {value = dense<-3.0> : tensor<1x2x4xf32>} : () -> tensor<1x2x4xf32>

      %%1 = "tf.%s"(%%arg0, %%arg1) {T = f32, adj_x = false, adj_y = false, grad_x = false, grad_y = false, device = ""} : (tensor<1x4x2xf32>, tensor<1x2x4xf32>) -> tensor<1x4x4xf32>

      func.return %%1 : tensor<1x4x4xf32>
    }
  })";
  std::string mat_mul_method =

    slice_sizes = #vhlo.tensor_v1<dense<[1, 2, 2]> : tensor<3xi64>>,
    indices_are_sorted = #vhlo.bool_v1<false>
  }> : (tensor<3x4x2xi32>, tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32>
  return %result : tensor<2x3x2x2xi32>
}


// CHECK: func.func private @gather(%arg0: tensor<3x4x2xi32>, %arg1: tensor<2x3x2xi64>) -> tensor<2x3x2x2xi32> {

func.func @valid_unranked_inputs_on_reshape(%arg0: tensor<3x4xi32>, %arg1: tensor<*xi32>) -> tensor<3x4xi32> {
  // CHECK: "tfl.reshape"(%arg0, %arg1)
  %0 = "tfl.reshape"(%arg0, %arg1) : (tensor<3x4xi32>, tensor<*xi32>) -> tensor<3x4xi32>
  func.return %0 : tensor<3x4xi32>
}

// -----

// CHECK-LABEL: valid_one_dynamic_dim_on_reshape
func.func @valid_one_dynamic_dim_on_reshape(%arg0: tensor<3x4xi32>) -> tensor<1x3x4xi32> {

}

// -----

// CHECK-LABEL: tensorlistConcat
func.func @tensorlistConcat(%arg0: tensor<3x2x2xf32>, %lead: tensor<i64>) -> (tensor<?x2xf32>, tensor<0xi64>) {
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %list = "tf.TensorListFromTensor"(%arg0, %cst) : (tensor<3x2x2xf32>, tensor<2xi32>) -> tensor<!tf_type.variant<tensor<f32>>>

    // CHECK:      metadata
    // CHECK-SAME: args
    // CHECK-SAME: shape
    // CHECK-SAME: dim
    // CHECK-SAME: size: 1
    // CHECK-SAME: dim
    // CHECK-SAME: size: 2
    // CHECK-SAME: dim
    // CHECK-SAME: size: 3

    func.return %0: tensor<1x2x3xi32>
  }