// CHECK:           %[[VAL_3:.*]] = "tf.LeftShift"(%[[VAL_1]], %[[VAL_0]]) : (tensor<4xi32>, tensor<1xi32>) -> tensor<4xi32>
// CHECK:           return %[[VAL_2]], %[[VAL_3]] : tensor<4xi32>, tensor<4xi32>
// CHECK:         }
func.func @broadcast_shift_left(%arg0: tensor<1xi32>, %arg1: tensor<4xi32>) -> (tensor<4xi32>, tensor<4xi32>) {

  %s0 = "tf.Const"() {value = dense<[4, 4]> : tensor<2xi32>} : () -> tensor<2xi32>
  %s1 = "tf.Const"() {value = dense<[1, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
  // expected-error @+1 {{requires dimension 0 size of 'r1' to be 2 but got 3}}
  %r0, %r1 = "tf.BroadcastGradientArgs"(%s0, %s1) : (tensor<2xi32>, tensor<2xi32>) -> (tensor<0xi32>, tensor<3xi32>)
  func.return %r0, %r1 : tensor<0xi32>, tensor<3xi32>
}

// CHECK-NEXT:      %10 = "tfl.concatenation"(%9, %7, %8) <{axis = 0 : i32, fused_activation_function = "NONE"}> : (tensor<1xi32>, tensor<1xi32>, tensor<1xi32>) -> tensor<3xi32>
// CHECK-NEXT:      %11 = mhlo.dynamic_reshape %arg0, %10 : (tensor<2x2x?x3xf32>, tensor<3xi32>) -> tensor<2x?x3xf32>
// CHECK-NEXT:      %12 = "tfl.shape"(%arg1) : (tensor<2x4x?x3xf32>) -> tensor<4xi32>

  %0 = "tfr.cast"(%arg0) : (tensor<*xf32>) -> !tfr.tensor
  %1 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<*xi32>
  %2 = "tfr.cast"(%0) : (!tfr.tensor) -> tensor<2xi32>
  func.return %1, %2 : tensor<*xi32>, tensor<2xi32>

// CHECK: %[[tf_cast_unranked:.*]] = "tf.Cast"(%arg0) <{Truncate = false}> : (tensor<*xf32>) -> tensor<*xi32>

  func.return %0#0 : tensor<2xi32>
}

// -----

func.func @unpack(%arg0: tensor<2x3xi32>) -> tensor<3xi32> {
  // CHECK: "tfl.unpack"(%arg0) <{axis = -2 : i32, num = 2 : i32}>
  %0:2 = "tfl.unpack"(%arg0) {axis = -2 : i32, num = 2 : i32} : (tensor<2x3xi32>) -> (tensor<3xi32>, tensor<3xi32>)
  func.return %0#0 : tensor<3xi32>
}

  %cst_perm = arith.constant dense<0> : tensor<1xi32>

  // CHECK: %[[CST:.*]] = arith.constant dense<{{\[}}1, 2, 3]> : tensor<3xi32>
  // CHECK: return %[[CST]]
  %0 = "tfl.transpose"(%cst, %cst_perm) : (tensor<3xi32>, tensor<1xi32>) -> tensor<3xi32>
  func.return %0 : tensor<3xi32>
}

// CHECK-LABEL: @transpose_dynamic
func.func @transpose_dynamic() -> tensor<?xi32> {
  %cst = arith.constant dense<[1, 2, 3]> : tensor<3xi32>

// CHECK-NEXT: %[[s3:.*]] = "tfl.shape"(%arg1) : (tensor<?x?x4x!quant.uniform<u8:f32, 1.000000e+00>>) -> tensor<3xi32>
// CHECK-NEXT: %[[s2]], %[[s3]] : tensor<?xi32>, tensor<3xi32>
}

// CHECK-LABEL: QuantizeMultipleUsers
func.func @QuantizeMultipleUsers(%arg1: tensor<?x!quant.uniform<u8:f32, 1.0>>) -> (tensor<1xi32>,tensor<1xi32>) {

  %0 = "tfl.pseudo_const"() {value = dense<0> : tensor<3xi32>} : () -> tensor<3xi32>
  %1 = "tfl.pseudo_const"() {value = dense<[1, 384, 32]> : tensor<3xi32>} : () -> tensor<3xi32>
  %2 = func.call @simple_test(%arg0, %0, %1) {tac.interface_name = "func1"} : (tensor<4x384x32xf32>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x384x32xf32>
  func.return %2 : tensor<1x384x32xf32>
}

// PARTIAL-LABEL: @simple_test

#loc3])\0Amodule {\0A  func.func @host_func(%arg0: tensor<3360x?xi32> loc(fused[#loc1, #loc2, #loc3]), %arg1: tensor<3xi32> loc(fused[#loc1, #loc2, #loc3]), %arg2: tensor<i32> loc(fused[#loc1, #loc2, #loc3]), %arg3: tensor<2xi32> loc(fused[#loc1, #loc2, #loc3])) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>) {\0A    %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) -> tensor<3x1120x?xi32> loc(#loc9)\0A    %1:3 = ...

  // CHECK-DAG: [[OFFSET_2:%.*]] = "tf.Const{{.*}} dense<[0, 5, 0]> : tensor<3xi32>

  %offset:3 = "tf.ConcatOffset"(%concat_dim, %shape0, %shape1, %shape2) : (tensor<i32>, tensor<3xi32>, tensor<3xi32>, tensor<3xi32>) -> (tensor<3xi32>, tensor<3xi32>, tensor<3xi32>)

  // CHECK: return [[OFFSET_0]], [[OFFSET_1]], [[OFFSET_2]]
  func.return %offset#0, %offset#1, %offset#2: tensor<3xi32>, tensor<3xi32>, tensor<3xi32>
}