// dilations, etc...). This test only verifies that changing convolution data
// layout will update all the attributes.

// CHECK-LABEL: func @transposeConv2D
func.func @transposeConv2D(%input: tensor<1x3x32x32xf32>, %filter: tensor<1x1x3x8xf32>) -> tensor<1x8x7x6xf32> {

  // CHECK: %[[ARG_PERM:.*]] = "tf.Const"() <{value = dense<[0, 2, 3, 1]> : tensor<4xi64>}>
  // CHECK: %[[ARG_TRANSPOSE:[0-9]*]] = "tf.Transpose"(%arg0, %[[ARG_PERM]])

  // CHECK-LABEL: testDilatedConv1DExpandH
  // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x3xf32>, [[FILTER:%.*]]: tensor<1x5x3x8xf32>)
  // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<-3> : tensor<i32>}> : () -> tensor<i32>
  // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x3xf32>, tensor<i32>) -> tensor<1x1x128x3xf32>

      %6 = "tf.Identity"(%5) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>
      "tf.Yield"(%4, %6) {device = ""} : (tensor<i1>, tensor<1x3xf32>) -> ()
    }) {_lower_using_switch_merge = true, _read_only_resource_inputs = [], device = ""} : (tensor<i1>) -> (tensor<i1>, tensor<1x3xf32>)
    %3 = "tf.Identity"(%2#1) {device = ""} : (tensor<1x3xf32>) -> tensor<1x3xf32>

//CHECK:func.func private @slice(%arg0: tensor<160x20x1xf32>) -> tensor<1x1x1xf32> {
//CHECK-NEXT: %0 = "vhlo.slice_v1"(%arg0) <{limit_indices = #vhlo.tensor_v1<dense<0> : tensor<3xi64>>, start_indices = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>, strides = #vhlo.tensor_v1<dense<1> : tensor<3xi64>>}> : (tensor<160x20x1xf32>) -> tensor<1x1x1xf32> 
//CHECK-NEXT: return %0 : tensor<1x1x1xf32>
//CHECK-NEXT:}

  %0 = "tfl.var_handle"() {container = "", shared_name = "states"} : () -> tensor<!tf_type.resource<tensor<1x2x3xf32>>>
  %1 = "tfl.read_variable"(%0) : (tensor<!tf_type.resource<tensor<1x2x3xf32>>>) -> tensor<1x2x3xf32>
  %2 = "tfl.quantize"(%1) {qtype = tensor<1x2x3x!quant.uniform<u8<1:255>:f32, 0.047244094488188976:128>>} : (tensor<1x2x3xf32>) -> tensor<1x2x3x!quant.uniform<u8<1:255>:f32, 0.047244094488188976:128>>

    func.func @main(%a : tensor<5x14x1xf32>, %b : tensor<1x14x32xf32>) -> tensor<?x?x?xf32> attributes {tf.entry_function = {control_outputs = "", inputs = "i,j", outputs = "identity_RetVal"}} {
      %c = "mhlo.einsum"(%a, %b) {einsum_config = "bji,bjk->bik"} : (tensor<5x14x1xf32>, tensor<1x14x32xf32>) -> tensor<?x?x?xf32>
      return %c : tensor<?x?x?xf32>
    }
  })";

  // CHECK: [[Result:%.*]] = "tf.Case"(%arg0, %arg1)
  // CHECK-SAME: branches = [@tf.CaseRegion_branch0{{.*}}, @tf.CaseRegion_branch1{{.*}}]
  // CHECK-SAME: is_stateless = false
  %0 = "tf.CaseRegion"(%arg0) ({
    %1 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<!tf_type.resource<tensor<1x2x3xf32>>>) -> tensor<!tf_type.resource>

func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32> {
  %0 = "tf.Identity"(%arg0) : (tensor<1x3xf32>) -> tensor<1x3xf32>
  func.return %0 : tensor<1x3xf32>
}

// -----

func.func @test_batch_function_with_invalid_symbol(%arg0: tensor<1x3xf32>, %arg1: tensor<!tf_type.resource<tensor<1x3xf32>>>) -> () {

  // CHECK: return %[[RESHAPE]] : tensor<1x2x3xf32>
}

// CHECK-LABEL: testConvertPackToReshapeAxis1
func.func @testConvertPackToReshapeAxis1(%arg0: tensor<2x3xf32>) -> tensor<2x1x3xf32> {
  %0 = "tf.Pack"(%arg0) {axis = 1 : i64, device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<2x3xf32>) -> tensor<2x1x3xf32>
  func.return %0 : tensor<2x1x3xf32>

func.func @QuantizedCatsAddRequantsTest(%arg0: tensor<1x1xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xf32>, %arg3: tensor<1x1xf32>) -> (tensor<1x4xf32>, tensor<1x3xf32>) {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[-0.440728068, 0.189515018]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>
  %1 = "quantfork.stats"(%arg1) {layerStats = dense<[-0.154693216, 0.26483655]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>