// 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>

// Ensure bucketize roundtrip exactly

func.func @main(%arg0: tensor<3x2xf32>) -> tensor<3x2xi32> {
  // CHECK-LABEL: @main
  // CHECK: "tfl.bucketize"(%arg0) <{boundaries = [0.000000e+00 : f32, 1.000000e+01 : f32, 1.000000e+02 : f32]}> : (tensor<3x2xf32>) -> tensor<3x2xi32>
  %0 = "tfl.bucketize"(%arg0) {boundaries = [0.0 : f32, 10.0 : f32, 100.0 : f32]} : (tensor<3x2xf32>) -> tensor<3x2xi32>

// Ensure "quantfork.stats" roundtrip exactly

func.func @main(%arg0: tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32> {
// CHECK-LABEL: @main
// CHECK: %[[RES0:.*]] = "quantfork.stats"(%arg0) <{layerStats = dense<[0.000000e+00, 2.550000e+02]> : tensor<2xf32>}> : (tensor<1x512x672x8xf32>) -> tensor<1x512x672x8xf32>

// CHECK-LABEL: func @device_test
func.func @device_test(
    %arg0: tensor<3x1xf32> {tf_saved_model.index_path = [0]},
    %arg1: tensor<1x3xf32> {tf_saved_model.index_path = [0]})
      -> (tensor<3x3xf32> {tf_saved_model.index_path = []}) {
  // CHECK: {{%.*}} = corert.get_op_handler %arg0 "/device:GPU:0"

// RUN: stablehlo-quant-opt %s -split-input-file -verify-diagnostics \
// RUN:   -stablehlo-test-pre-calibration-component | FileCheck %s

func.func @main(%arg0: tensor<1x4xf32>) -> tensor<1x3xf32> {
  %0 = stablehlo.constant dense<1.0> : tensor<4x3xf32>
  %1 = stablehlo.dot_general %arg0, %0, contracting_dims = [1] x [0], precision = [DEFAULT, DEFAULT] : (tensor<1x4xf32>, tensor<4x3xf32>) -> tensor<1x3xf32>
  return %1 : tensor<1x3xf32>
}

    %1 = "tf.Less"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i1>
    func.return %1 : tensor<i1>
  }

  func.func @while_body_add2(%arg0: tensor<i32>) -> tensor<i32> {
    %0 = "tf.Const"() {device = "/device:CPU:0", value = dense<2> : tensor<i32>} : () -> tensor<i32>
    %1 = "tf.Add"(%arg0, %0) {device = "/device:CPU:0"} : (tensor<i32>, tensor<i32>) -> tensor<i32>

  // CHECK: return

  func.func private @composite_dot_general_with_relu_fn_1(%arg0: tensor<10x1x1024xf32>, %arg1: tensor<10x1024x3xf32>) -> tensor<10x1x3xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
    %cst = stablehlo.constant dense<0.000000e+00> : tensor<10x1x3xf32>

module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 1660 : i32}} {
  func.func @main(%arg0: tensor<2x3xi32>) -> tensor<2x3xi32> attributes {tf.entry_function = {control_outputs = "", inputs = "args_tf_0", outputs = "Identity"}} {
    %0 = tf_executor.graph {
      %outputs, %control = tf_executor.island wraps "tf.Identity"(%arg0) {device = ""} : (tensor<2x3xi32>) -> tensor<2x3xi32>

        module = ir.Module.parse(module_bytecode)
        self.assertIn('stablehlo.add %arg0, %cst : tensor<4xf32>', str(module))

  def test_tf_mlir_to_stablehlo(self):
    assembly = """
      module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
        func.func @main(%arg0 : tensor<?xf32>) -> tensor<?xf32> {