// CHECK-LABEL: func @quantized_types
func.func @quantized_types(%arg0: tensor<!tf_type.resource<tensor<1x3x!tf_type.quint8>>>,
                      %arg1: tensor<!tf_type.resource<tensor<1x3x!tf_type.quint16>>>,
                      %arg2: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint8>>>,
                      %arg3: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint16>>>,
                      %arg4: tensor<!tf_type.resource<tensor<1x3x!tf_type.qint32>>>) {

// Ensure basic_lstm roundtrip exactly

func.func @main(%arg0: tensor<1x384xf32>, %arg1: tensor<1x96xf32>, %arg2: tensor<384x480xf32>, %arg3: tensor<384xf32>, %arg4: tensor<1x96xf32>) -> tensor<1x96xf32> {
// CHECK-LABEL: @main

  // CHECK:      %[[READ0:.*]] = "tf.ReadVariableOp"(%arg0)
  // CHECK-NEXT: %[[READ2:.*]] = "tf.ReadVariableOp"(%arg2)
  // CHECK-NEXT: tf_device.replicate
  // CHECK-NEXT:   %[[READ1:.*]] = "tf.ReadVariableOp"(%arg1)
  // CHECK-NEXT:   "tf.OpA"(%[[READ0]], %[[READ1]], %[[READ2]])
  tf_device.replicate {n = 2 : i32} {
    %0 = "tf.ReadVariableOp"(%arg0) : (tensor<*x!tf_type.resource<tensor<f32>>>) -> tensor<f32>

  }

  func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
    return %0 : tensor<1x3xf32>
  }
}

// CHECK-LABEL: quantize_dot_general_per_tensor
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>

  }

  func.func private @composite_dot_general_fn(%arg0: tensor<1x2xf32>, %arg1: tensor<2x3xf32>) -> tensor<1x3xf32> attributes {_from_xla_call_module} {
    %0 = stablehlo.dot_general %arg0, %arg1, contracting_dims = [1] x [0] : (tensor<1x2xf32>, tensor<2x3xf32>) -> tensor<1x3xf32>
    return %0 : tensor<1x3xf32>
  }
}

// CHECK-LABEL: quantize_dot_general_fn
// CHECK-SAME: %[[ARG0:.+]]: tensor<1x2xf32>

}

// CHECK-LABEL: func @identity_n
// CHECK-SAME:    (%[[ARG0:.*]]: tensor<i32>, %[[ARG1:.*]]: tensor<f32>)
func.func @identity_n(%arg0: tensor<i32>, %arg1: tensor<f32>) -> (tensor<i32>, tensor<f32>) {
  // CHECK-NOT: "tf.IdentityN"
  %0:2 = "tf.IdentityN"(%arg0, %arg1) : (tensor<i32>, tensor<f32>) -> (tensor<i32>, tensor<f32>)
  // CHECK: return %[[ARG0]], %[[ARG1]]
  func.return %0#0, %0#1 : tensor<i32>, tensor<f32>

module {
func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>, %arg2: tensor<1xf32>, %arg3: tensor<1xf32>) -> tensor<2x1xf32> attributes {tf.entry_function = {inputs = "input0,input1,input2,input3", outputs = "output"}} {
  %0 = "tfl.add"(%arg0, %arg1) {fused_activation_function = "RELU6"} : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>

func.func @main(%arg0: tensor<4x10x15xf32>, %arg1: tensor<4x15x17xf32>) -> tensor<4x10x17xf32> {
  %0 = "tfl.batch_matmul"(%arg0, %arg1) {adj_x = false, adj_y = false} : (tensor<4x10x15xf32>, tensor<4x15x17xf32>) -> tensor<4x10x17xf32>
  func.return %0:  tensor<4x10x17xf32>

// CHECK-LABEL: main

// Identity function to make the exporter happy
func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> {
  func.return %arg0 : tensor<4xi8>
}

//CHECK:func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> attributes {tf.entry_function = {inputs = "arg0", outputs = "arg0"}} {
//CHECK: return %arg0 : tensor<4xi8>
//CHECK:}

func.func @custom_tf_op(%arg0: tensor<1x320x1x1xf32>, %arg1: tensor<2xi32>) -> tensor<1x1600x1x1xf32> {

  func.func private @composite_conv2d_with_bias_and_relu6_fn_10(%arg0: tensor<1x3x2x3xf32>, %arg1: tensor<2x3x3x2xf32>, %arg2: tensor<2xf32>) -> tensor<1x2x2x2xf32> attributes {tf.tf_quant.composite_function} {
    %0 = "quantfork.stats"(%arg1) {layerStats = dense<[-3.54062747, 0.54742622]> : tensor<2xf32>} : (tensor<2x3x3x2xf32>) -> tensor<2x3x3x2xf32>