// CHECK:   %[[MUL_0:.*]] = "tf.Mul"(%arg1, %arg0) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
// CHECK:   return %[[MUL_0]] : tensor<1xf32>
// CHECK: }

// CHECK:         }
func.func @tfl_wrapped_jax_random_normal(%arg0: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
  // This is a fake jax random normal body.
  %0 = stablehlo.constant dense<0.0> : tensor<12xf32>
  %1 = "stablehlo.reshape"(%0) : (tensor<12xf32>) -> tensor<3x4xf32>
  %2 = "stablehlo.tuple"(%1) : (tensor<3x4xf32>) -> tuple<tensor<3x4xf32>>
  func.return %2 : tuple<tensor<3x4xf32>>
}

func.func @listPopBack(%arg0: tensor<!tf_type.variant<tensor<2xi32>>>, %arg1: tensor<1xi32>) -> (tensor<!tf_type.variant<tensor<2xi32>>>, tensor<2xi32>)  {
  %0, %1 = "tf.TensorListPopBack"(%arg0, %arg1) : (tensor<!tf_type.variant<tensor<2xi32>>>, tensor<1xi32>) -> (tensor<!tf_type.variant<tensor<2xi32>>>, tensor<2xi32>)

module {
  func.func @same_scale_test(%arg0: tensor<*xf32>) -> tensor<*xf32> {
    %cst = arith.constant dense<[-1, 144]> : tensor<2xi32>
    %cst_1 = arith.constant dense<1.0> : tensor<144x10xf32>
    %cst_2 = arith.constant dense<0.1> : tensor<10xf32>
    %0 = "quantfork.qcast"(%arg0) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.05:-10>>

// CHECK: %[[v5:.*]] = "tf.Gather"(%3, %[[cst_1]]) <{validate_indices = true}> : (tensor<3xi32>, tensor<2xi32>) -> tensor<2xi32>
// CHECK: %[[v6:.*]] = "tf.Gather"(%4, %[[cst_2]]) <{validate_indices = true}> : (tensor<4xi32>, tensor<1xi32>) -> tensor<1xi32>
// CHECK: %[[v7:.*]] = "tf.Concat"(%[[cst_3]], %5, %6) : (tensor<i32>, tensor<2xi32>, tensor<1xi32>) -> tensor<3xi32>

    "tf.AssignVariableOp"(%arg0, %3) <{validate_shape = false}> : (tensor<*x!tf_type.resource<tensor<1xf32>>>, tensor<1xf32>) -> ()
    tf_device.return
  }
  func.return
}

func.func private @write_chain_func(%arg0: tensor<1xf32>) -> (tensor<1xf32>) {
  %cst = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>

// PerTensor-LABEL: QuantizeGatherWeightOnly
func.func @QuantizeGatherWeightOnly(%arg0: tensor<3xi32>) -> tensor<3x3x3x3xf32> {
  %w = arith.constant dense<1.270000e+02> : tensor<64x3x3x3xf32>
  %emb = "tfl.gather"(%w, %arg0) {axis = 0 : i32, batch_dims = 0 : i32} : (tensor<64x3x3x3xf32>, tensor<3xi32>) -> tensor<3x3x3x3xf32>

  %size = "tf.Const"() {device = "/device:CPU:0", value = dense<1> : tensor<i32>} : () -> (tensor<i32>)
  %indices = "tf.Const"() {device = "/device:CPU:0", value = dense<[0]> : tensor<1xi32>} : () -> (tensor<1xi32>)
  // CHECK: tfrt_fallback_async.executeop key({{.*}}) cost({{.*}}) device("/job:localhost/replica:0/task:0/device:CPU:0") "tf.TensorArrayV3"

// PerChannel-DAG: %[[scale2:.*]] = "tf.Const"() <{value = dense<{{[0-9\.Ee\+\-]+}}> : tensor<6xf32>}> : () -> tensor<6xf32>
// PerChannel-DAG: %[[zp1:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<3xi32>}> : () -> tensor<3xi32>
// PerChannel-DAG: %[[zp2:.*]] = "tf.Const"() <{value = dense<{{[0-9]+}}> : tensor<6xi32>}> : () -> tensor<6xi32>

  %3:2 = "tf.TPUReplicatedOutput"(%2#0) : (tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)
  %5:2 = "tf.TPUPartitionedOutputV2"(%3#0) {_XlaSharding = "", partition_dims = []} : (tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)
  %7:2 = "tf.TPUPartitionedOutputV2"(%3#1) {_XlaSharding = "", partition_dims = []} : (tensor<?xi32>) -> (tensor<?xi32>, tensor<?xi32>)