// CHECK-NEXT: return %0 : tensor<2x3x2x2xi32>
// CHECK-NEXT:}

func.func @transpose(%arg0: tensor<2x3x2xi32>) -> tensor<2x3x2xi32> {
  %0 = "vhlo.transpose_v1"(%arg0) <{permutation = #vhlo.tensor_v1<dense<[2, 1, 0]> : tensor<3xi64>>}> : (tensor<2x3x2xi32>) -> tensor<2x3x2xi32>
  return %0 : tensor<2x3x2xi32>
}

    %cst = "tf.Const"() {value = dense<1> : tensor<2x3x2xi32>} : () -> tensor<2x3x2xi32>
    %0 = "quantfork.stats"(%arg) {layerStats = dense<[4.00000000e-6, 9.80000000e-1]> : tensor<2xf32>} : (tensor<3x4x2xf32>) -> tensor<3x4x2xf32>

// -----

// Tests that int constants are not folded.

// CHECK-LABEL: transpose_int
func.func @transpose_int() -> tensor<3x2xi32> {
  %0 = stablehlo.constant dense<0> : tensor<2x3xi32>
  %1 = stablehlo.transpose %0, dims = [1, 0] : (tensor<2x3xi32>) -> tensor<3x2xi32>
  return %1 : tensor<3x2xi32>
}
// CHECK: transpose

// -----

// Tests that transposing an argument cannot be folded.

// CHECK-LABEL: transpose_arg

    %16 = "tf.Cast"(%15) {Truncate = false} : (tensor<2x2x2xi32>) -> tensor<2x2x2xf32>
    %17 = "tf.Mul"(%16, %cst_0) : (tensor<2x2x2xf32>, tensor<f32>) -> tensor<2x2x2xf32>
    %18 = "tf.AddV2"(%17, %cst) : (tensor<2x2x2xf32>, tensor<f32>) -> tensor<2x2x2xf32>
    %19 = "tf.Floor"(%18) : (tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
    %20 = "tf.ClipByValue"(%19, %cst_5, %cst_6) : (tensor<2x2x2xf32>, tensor<f32>, tensor<f32>) -> tensor<2x2x2xf32>

  // CHECK: return %[[RESULT]]

  %indices = "tf.Const"() {value = dense<[1, 0]> : tensor<2xi32>} : () -> tensor<2xi32>
  %0 = "tf.DynamicStitch"(%indices, %arg0) : (tensor<2xi32>, tensor<2x2x2xf32>) -> tensor<2x2x2xf32>
  func.return %0 : tensor<2x2x2xf32>
}

// CHECK-LABEL: func @DynamicStitch_dynamic

    %7 = stablehlo.broadcast_in_dim %6, dims = [2, 1] : (tensor<1x3xf32>) -> tensor<2x3x2xf32>
    %8 = "quantfork.qcast"(%7) {volatile} : (tensor<2x3x2xf32>) -> tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>
    %9 = "quantfork.dcast"(%8) : (tensor<2x3x2x!quant.uniform<i8:f32, 0.13170163023705575:-1>>) -> tensor<2x3x2xf32>
    return %9 : tensor<2x3x2xf32>
  }

  // CHECK: quantized_dot_general_fn_1

  %2 = "tfl.add"(%cst_0, %cst_2) {fused_activation_function = "NONE"} : (tensor<    2xi32>, tensor<2x2x2xi32>) -> tensor<2x2x2xi32>

  func.return %0, %1, %2 : tensor<2x2xi32>, tensor<2x2x2xi32>, tensor<2x2x2xi32>

// CHECK-DAG:  %[[CST:.*]] = arith.constant dense<{{\[\[}}11, 22], [13, 24]]> : tensor<2x2xi32>

  %0 = "tfl.custom"(%arg0) {custom_code = "stablehlo.broadcast_in_dim", custom_option = #tfl<const_bytes : "0x62726F6164636173745F64696D656E73696F6E73000201020119010101072C022401">} : (tensor<1x2xi32>) -> tensor<1x2x2xi32>
  func.return %0 : tensor<1x2x2xi32>
  }
}

// CHECK:      module {
// CHECK-NEXT:  func @main(%arg0: tensor<1x2xi32>) -> tensor<1x2x2xi32> {

    %1 = "tf.Const"() {value = dense<2> : tensor<2x2xi32>} : () -> tensor<2x2xi32>
    %2 = "tf.Identity"(%1) {device = ""} : (tensor<2x2xi32>) -> tensor<2x2xi32>
    // CHECK: SpaceToBatchND{{.*}}-> tensor<4x98x130x128xf32>
    %3 = "tf.SpaceToBatchND"(%arg0, %0, %2) {device = ""} : (tensor<1x192x256x128xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x?x?x128xf32>
    func.return
  }

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

//CHECK: module attributes