// CHECK:   %{{.*}} = "tfl.dequantize"(%{{.*}}) : (tensor<1x401408x!quant.uniform<u8:f32, 3.906250e-03>>) -> tensor<1x401408xf32>

  %cst = arith.constant dense<[1, 401408]> : tensor<2xi32>

    func.return
}

func.func private @__reduce_func_1(%arg0: tensor<i64> {tf._user_specified_name = "args_0"}, %arg1: tensor<32xf32> {tf._user_specified_name = "args_1"}) -> (tensor<i64>) attributes {tf._tf_data_function = true, tf.signature.is_stateful} {
    %0 = "tf.JustPretend"(%arg1) : (tensor<32xf32>) -> (tensor<i64>)
    func.return %0 : tensor<i64>
}

// CHECK-LABEL: func @single_state_single_dataset_type_multiple_arguments

func.func private @simple_callee(%arg0: tensor<2xi32>) -> tensor<2xi32> {
  %graph = tf_executor.graph {
    tf_executor.fetch %arg0 : tensor<2xi32>
  }
  func.return %graph#0 : tensor<2xi32>
}

// Test with a TPUPartitionedCallOp that is not inlined.
func.func @dont_inline_tpu_partitioned_call(%arg0: tensor<2xi32> {tf._user_specified_name = "inputs_0"}, %arg1: tensor<2xi32> {tf._user_specified_name = "inputs_1"}) -> tensor<2xi32> {

  // CHECK-NEXT:  buffer: 6
  // CHECK-NEXT:  } ]
  // CHECK-NEXT:  signature_defs: [ ]
  // CHECK-NEXT:}

  %0 = "tfl.pseudo_const" () {value = dense<-1.23697901> : tensor<32xf32>} : () -> tensor<32xf32> loc("Const")

  // CHECK-NEXT:  buffer: 6
  // CHECK-NEXT:  } ]
  // CHECK-NEXT:  signature_defs: [ ]
  // CHECK-NEXT:}

  %0 = "tfl.pseudo_const" () {value = dense<-1.23697901> : tensor<32xf32>} : () -> tensor<32xf32> loc("Const")

    %4 = "tf.LayerNorm"(%a1, %a2, %a3, %a4) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>
   "tfl.yield"(%4) : (tensor<128x128xf32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable", device = ""} : (tensor<128x128xf32>, tensor<1xf32>, tensor<1xf32>, tensor<1xi32>) -> tensor<128x128xf32>

func.func @bias_add(%arg0: tensor<1x10x10x32xf32>, %arg1: tensor<32xf32>) -> tensor<1x10x10x32xf32> {
  %0 = "tf.BiasAdd"(%arg0, %arg1) {T = "tfdtype$DT_FLOAT", data_format = "NHWC"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32>
  func.return %0 : tensor<1x10x10x32xf32>
// CHECK: %[[BIASADD_0:.*]] = "tf.BiasAdd"(%arg0, %arg1) <{data_format = "NHWC"}> {T = "tfdtype$DT_FLOAT"} : (tensor<1x10x10x32xf32>, tensor<32xf32>) -> tensor<1x10x10x32xf32>

  %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>)

  attr {
    key: "dtype"
    value {
      type: DT_INT32
    }
  }
}
versions {
  producer: 27
}

# CHECK-LABEL: func @main
# CHECK-SAME:  (%[[ARG_0:[a-z0-9]+]]: tensor<10xi32>, %[[ARG_1:[a-z0-9]+]]: tensor<10xi32>) -> tensor<*xi32>
# CHECK-SAME:  control_outputs = ""
# CHECK-SAME:  inputs = "input0,input1"
# CHECK-SAME:  outputs = "Sub"

    // CHECK: [[TENSOR:%.+]] = "tf.TensorScatterAdd"([[READ]], [[EXPAND]], [[UPDATE]]) : (tensor<*xi32>, tensor<?x1xi32>, tensor<?x?x?xi32>) -> tensor<*xi32>
    // CHECK: "tf.AssignVariableOp"([[VAR]], [[TENSOR]])
    "tf.ResourceScatterAdd"(%resource, %indices, %updates) : (tensor<*x!tf_type.resource<tensor<*xi32>>>, tensor<?xi32>, tensor<?x?x?xi32>) -> ()

    tf_device.return