%0:4 = "tfl.split"(%arg0, %arg1) {num_splits = 4 : i32, tac.device = "CPU"} : (tensor<i32>, tensor<1x8x8x1024xf32>) -> (tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>)
  func.return %0#0, %0#1, %0#3 : tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>, tensor<1x8x8x256xf32>
}

  %1 = "tfl.mul"(%0, %cst_2) {fused_activation_function = "NONE"} : (tensor<1x8x8x256xf32>, tensor<1x1x1x256xf32>) -> tensor<1x8x8x256xf32>
  func.return %1 : tensor<1x8x8x256xf32>

// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<2.800000e+00> : tensor<256x3x3x207xf32>

  return %0 : tensor<1x1x1x2xf32>
}
func.func private @XlaCallModule_aten.avg_pool2d.default.impl_5(%arg0: tensor<1x1x1x7xf32>) -> tensor<1x1x1x2xf32>

// CHECK-LABEL: avg_pool2d_6
// CHECK: %cst = arith.constant dense<[0, 2, 3, 1]> : tensor<4xi32>

// test stablehlo roundtrip

module {
func.func @main(%arg0: tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32> {
  %0 = "vhlo.logistic_v1"(%arg0) : (tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32>
  %1 = "tfl.exp"(%0) : (tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32> loc("exp")
  func.return %1 : tensor<1x1x1x96xf32>
}
}

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

func.func @exp(%arg0: tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32> {
  %0 = "vhlo.exponential_v1" (%arg0) : (tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32>
  func.return %0 : tensor<1x1x1x96xf32>
}

//CHECK:func.func private @exp(%arg0: tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32> {
//CHECK-NEXT: %0 = "vhlo.exponential_v1"(%arg0) : (tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32>
//CHECK-NEXT: return %0 : tensor<1x1x1x96xf32>
//CHECK-NEXT:}

module attributes {tfl.metadata = {"keep_stablehlo_constant" = "true"}} {
 func.func @main () -> tensor<1x1x1x96xf32> {
  %0 = "vhlo.constant_v1"() <{value = #vhlo.tensor_v1<dense<0.000000e+00> : tensor<f32>>}> : () -> tensor<1x1x1x96xf32>
  func.return %0 : tensor<1x1x1x96xf32>
 }
}

//CHECK: func.func @main() -> tensor<1x1x1x96xf32> attributes {tf.entry_function = {outputs = "vhlo.constant_v1"}} {

func.func @main(tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32> {
^bb0(%arg0: tensor<1x6x6x16xf32>):
  // CHECK: "tfl.average_pool_2d"(%{{.*}}) <{filter_height = 3 : i32, filter_width = 6 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 3 : i32, stride_w = 1 : i32}> : (tensor<1x6x6x16xf32>) -> tensor<1x1x1x16xf32>

// CHECK:  [[VAL_3:%.*]] = "tfl.reshape"([[VAL_2]], %{{.*}}) : (tensor<1x1x1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>

  func.return %1 : tensor<1x1x1x16xf32>

// CHECK: %[[avgp:.*]] = "tfl.average_pool_2d"(%arg0)
// CHECK: %[[dq:.*]] = "tfl.dequantize"(%[[avgp]]) : (tensor<1x1x1x16x!quant.uniform<u8:f32, 7.812500e-03:128>>) -> tensor<1x1x1x16xf32>
// CHECK: return %[[dq]] : tensor<1x1x1x16xf32>
}

// CHECK-LABEL: QuantizeReshape2D

    %cst = "tf.Const"() {value = dense<3.00000000e-1> : tensor<2x3x3x2xf32>} : () -> tensor<2x3x3x2xf32>
    %cst_0 = "tf.Const"() {value = dense<4.00000000e-1> : tensor<1x1x1x2xf32>} : () -> tensor<1x1x1x2xf32>
    %0 = "quantfork.stats"(%arg0) {layerStats = dense<[6.00000000e-6, 9.00000000e-1]> : tensor<2xf32>} : (tensor<1x3x4x3xf32>) -> tensor<1x3x4x3xf32>
    %1 = "tf.XlaCallModule"(%0, %cst, %cst_0) {