// CHECK: return

// -----

// No CustomAggregator ops exist.
func.func private @composite_conv2d_with_bias_and_relu6_fn_1(%arg0: tensor<1x3x4x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x2x2x2xf32> attributes {tf_quant.composite_function} {

                << context->DebugString(handle);
      }
    }

    int index = -1;
    if (n->type_string() == "_Arg") {
      // NOTE: during runtime, Placeholder ops will be replaced as `_Arg` ops.
      // And Args must have `index` attribute.
      TF_RETURN_IF_ERROR(GetNodeAttr(n->attrs(), "index", &index));
    } else if (n->type_string() == "Placeholder") {

func.func @QuantizeAssignVariableWithDequantAndEqualType(%arg0 : tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>) -> tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>> {
  %0 = "tfl.var_handle"() : () -> tensor<!tf_type.resource>
  %1 = "tfl.dequantize"(%arg0) : (tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>) -> tensor<1x2x1x3xf32>
  "tfl.assign_variable"(%0, %1) : (tensor<!tf_type.resource>, tensor<1x2x1x3xf32>) -> ()
  func.return %arg0 : tensor<1x2x1x3x!quant.uniform<i8:f32, 1.0>>

module attributes {tf_saved_model.semantics} {
  func.func @no_call_once(%arg0: tensor<i64> {tf_saved_model.index_path = ["x"]}) -> (tensor<i64> {tf_saved_model.index_path = ["r"]})
  attributes {tf.entry_function = {control_outputs = "", inputs = "input:0", outputs = "output:0"}, tf_saved_model.exported_names = ["serving_default"]} {
    func.return %arg0 : tensor<i64>
    // CHECK-LABEL: no_call_once
    // CHECK-NOT: "tfl.call_once"
  }

module {
func.func @degenerateCase(%arg0: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tfl.add"(%arg0, %arg0) {tac.device = "GPU", fused_activation_function = "RELU6", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  func.return %0 : tensor<1xf32>
}
}

// CHECK:     func.func @degenerateCase(%arg0: tensor<1xf32>) -> tensor<1xf32> {

  %0:2 = "tfl.while"(%cst_0, %cst_1) ({
  ^bb0(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>):
    // CHECK: call @WhileOp_cond
    // CHECK-SAME: (tensor<*xi32>, tensor<*xf32>, tensor<i32>)
    %1 = func.call @WhileOp_cond(%arg0, %arg1, %cst) : (tensor<*xi32>, tensor<*xf32>, tensor<i32>) -> tensor<i1>
    "tfl.yield"(%1) : (tensor<i1>) -> ()
  },  {
  ^bb0(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>):
    // CHECK: call @WhileOp_body

        "Property<Boolean>"           | "boolean"        | "self.getProperty().getOrElse(false)"            | ".set(arg0)"       | []
        "Property<Long>"              | "long"           | "self.getProperty().getOrElse(0L)"               | ".set(arg0)"       | []

module {
func.func @tfInplaceUpdate(%arg0: tensor<2x1x2xf32>) -> tensor<2x1x2xf32> {
  %1 = arith.constant dense<1> : tensor<1xi32>
  %2 = arith.constant dense<2.0> : tensor<1x1x2xf32>
  %3 = "tf.InplaceUpdate"(%arg0, %1, %2) {device = ""}
    : (tensor<2x1x2xf32>, tensor<1xi32>, tensor<1x1x2xf32>) -> tensor<2x1x2xf32>
  func.return %3 : tensor<2x1x2xf32>
}

// RUN: tac-opt-all-backends -tfl-device-transform-nnapi %s -split-input-file -verify-diagnostics | FileCheck %s

func.func @mean_4d_keepdim(%arg0: tensor<1x48x48x512xf32>) -> tensor<1x1x1x512xf32> {
  %cst = arith.constant dense<[1, 2]> : tensor<2xi32>
  %0 = "tfl.mean"(%arg0, %cst) {keep_dims = true} : (tensor<1x48x48x512xf32>, tensor<2xi32>) -> tensor<1x1x1x512xf32>
  func.return %0 : tensor<1x1x1x512xf32>
}

func.func private @batched_function(%arg0: tensor<1x3xf32> {tf._user_specified_name = "0"}, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32> attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
  %0 = "tf.ReadVariableOp"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<1x3xf32>