// 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-NEXT:      %14 = "tfl.pseudo_const"() <{value = dense<[-1, -1, 0, -1]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT:      %15 = "tfl.pseudo_const"() <{value = dense<[-1, -1, -1, 0]> : tensor<4xi32>}> : () -> tensor<4xi32>
// CHECK-NEXT:      %16 = "tfl.pseudo_const"() <{value = dense<1> : tensor<i32>}> : () -> tensor<i32>

func.func @init() attributes {tf_saved_model.exported_names = ["__tf_saved_model_session_initializer_init"]} {
  %keys = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
  %values = "tf.Const"() {value = dense<[1, 2, 3, 4]> : tensor<4xi32> } : () -> tensor<4xi32>
  %handle = "tf.HashTableV2"() {container = "", key_dtype = i32, shared_name = "hash_table", use_node_name_sharing = false, value_dtype = i32} : () -> tensor<!tf_type.resource>

  func.return %0 : tensor<4xi16>
}

func.func @i32() -> tensor<4xi32> {
  // CHECK-LABEL: @i32
  // CHECK: value = dense<[1, 2, 3, 16909060]> : tensor<4xi32>
  // Check bytes come back in the right order
  %0 = "tfl.pseudo_const" () { value = dense<[1, 2, 3, 16909060]> : tensor<4xi32> } : () -> tensor<4xi32>
  func.return %0 : tensor<4xi32>
}

func.func @i64() -> tensor<4xi64> {
  // CHECK-LABEL: @i64

  %0 = "tfl.pseudo_const"() {value = dense<[1, 4, 384, 32]> : tensor<4xi32>} : () -> tensor<4xi32>
  // CHECK-NOT: tac.device tac.inference_type
  %1 = "tfl.pseudo_const"() {value = dense<[4, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
  // CHECK: tac.device = "CPU", tac.inference_type = "QUANTIZED_INT8"
  %2 = "tfl.reshape"(%arg0, %0) : (tensor<4x384x32x!quant.uniform<i8:f32, 0.2:-3>>, tensor<4xi32>) -> tensor<1x4x384x32x!quant.uniform<i8:f32, 0.2:-3>>

  %1 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>

  }

  return true;
}

// Matches %iota generated from the following code (rank 3 example):
//
// %iota_r1 = "mhlo.iota"(){iota_dimension = 0 : i32} : () -> tensor<44xi32>
// %iota = "mhlo.reshape"(%iota_r1): (tensor<44xi32>) -> tensor<1x1x44xi32>
//
// Where $dimensions is of size 1 and $dimensions[0] = 2.
//
// In general matches a 1-D Iota with multiple dimensions of size 1 added
// through a reshape.

    %cst = "tf.Const"() {value = dense<[0, 1, 2, 3]> : tensor<4xi32>} : () -> tensor<4xi32>
    %cst_0 = "tf.Const"() {value = dense<5.0> : tensor<f32>} : () -> tensor<f32>
    // Check that these constants are remained in constants.
    // CHECK: func private @while_cond
    // CHECK-DAG:  %[[CST:.*]]= "tf.Const"() <{value = dense<[0, 1, 2, 3]> : tensor<4xi32>}> : () -> tensor<4xi32>

  // CHECK-SAME: -> tensor<4xi32>
  // CHECK: %[[CLUSTER_RES:[0-9]*]]:2 = "tf_device.cluster"
  // CHECK: %[[COMPUTE_RES:[0-9]*]] = "tf.SomeComputation"(%[[RES_READ_VAL]]) : (tensor<4xi32>) -> tensor<4xi32>
  // CHECK: tf_device.return %[[COMPUTE_RES]], %[[COMPUTE_RES]]
  // CHECK-SAME: tensor<4xi32>, tensor<4xi32>
  // CHECK: {cluster_attr = "cluster_attr"}

}

func.func @dynamic_update_slice(%arg0: tensor<4x5xi32>, %arg1: tensor<1x5xi32>, %arg2: tensor<2xi32>) -> tensor<4x5xi32> {
  %0 = "tf.XlaDynamicUpdateSlice"(%arg0, %arg1, %arg2) : (tensor<4x5xi32>, tensor<1x5xi32>, tensor<2xi32>) -> tensor<4x5xi32>
  func.return %0 : tensor<4x5xi32>

// CHECK-LABEL:dynamic_update_slice
// CHECK: "tfl.dynamic_update_slice"(%arg0, %arg1, %arg2) : (tensor<4x5xi32>, tensor<1x5xi32>, tensor<2xi32>) -> tensor<4x5xi32>