# CHECK-DAG:       %[[VAL_3:.*]] = "tfl.no_value"() <{value}> : () -> none
# CHECK-DAG:       %[[VAL_6:.*]] = arith.constant dense<0> : tensor<i32>
# CHECK:           %[[VAL_7:.*]]:2 = "tfl.split"(%[[VAL_6]], %[[VAL_0]]) <{num_splits = 2 : i32}> : (tensor<i32>, tensor<2x5x3xf32>) -> (tensor<1x5x3xf32>, tensor<1x5x3xf32>)

      %cst = "tf.Const"() {value = dense<1024> : tensor<i32>} : () -> tensor<i32>
      %0:2 = "tf_device.launch"() ({
            %1 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>
            %2 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>

  %3 = "tfl.depthwise_conv_2d"(%arg0, %2, %0) {depth_multiplier = 4 : i32, dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 4 : i32, stride_w = 5 : i32} : (tensor<1x224x224x3xf32>, tensor<32x3x3x3xf32>, tensor<32xf32>) -> tensor<1x112x112x32xf32>

  func.return %3 : tensor<1x112x112x32xf32>

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {

// CHECK-LABEL: non_max_suppression_v4
func.func @non_max_suppression_v4(%arg0: tensor<3x4xf32>, %arg1: tensor<3xf32>, %arg2: tensor<f32>, %arg3: tensor<f32>) -> tensor<2xi32> {
  %max_size = mhlo.constant dense<2> : tensor<i32>
  // NO_FALLBACK: tf.NonMaxSuppressionV4
  // SUPPORTED_FALLBACK_DEVICE-NOT: tf.NonMaxSuppressionV4

// CHECK-LABEL: func @move_broadcast
func.func @move_broadcast(%arg0: tensor<f32>) -> () {
  // CHECK:      %[[ELEM_0:.*]] = "tf.Const"()
  // CHECK:      {ici_weight_distribution_mlir_bridge_marker = true}

// -----

// Tests that the converted `tf.Slice` has the correct number of dimensions
// when the output shape is known (`tensor<i32>` instead of `tensor<*xi32>`).

func.func @xla_gather_known_output_shape(%arg0: tensor<5xi32>, %arg1: tensor<1xi32>, %arg2: tensor<1xi32>) -> tensor<i32> {
  // dimension_numbers: {
  //   collapsed_slice_dims: 0
  //   start_index_map: 0
  // }

    %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
    func.return %0 : tensor<2x1xf32>
    // CHECK: %[[VAL_0:.*]] = arith.constant dense<[2, 1]> : tensor<2xi32>
    // CHECK: %[[CONCAT:.*]] = "tfl.concatenation"(%arg0, %arg1) <{axis = 0 : i32, fused_activation_function = "NONE"}> : (tensor<1xf32>, tensor<1xf32>) -> tensor<2xf32>

  // will be preserved since it has uses.
  %2 = "tf.MyCustomOp"(%1, %0) {fused_activation_function = "RELU", int_attr = 2 : i32}  : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
  // will be preserved since it has side-effect.
  "tf.MyCustomOp"(%1, %0) {fused_activation_function = "RELU", int_attr = 2 : i32}  : (tensor<4xf32>, tensor<4xf32>) -> tensor<4xf32>
  func.return %2 : tensor<4xf32>

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
  func.func @main() -> () {
    func.return
  }

# CHECK:  tf.TensorListReserve{{.*}}(tensor<2xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*xf32>>>

# Nested variant type.
# CHECK:  tf.TensorListReserve{{.*}}(tensor<2xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<*x!tf_type.variant>>>

# CHECK:  tf.TensorListSetItem{{.*}}(tensor<!tf_type.variant<tensor<*xf32>>>, tensor<i32>, tensor<2x2xf32>) -> tensor<*x!tf_type.variant>