func.func @transpose_simple_1d() -> tensor<2xf32> {
  %0 = stablehlo.constant dense<[0.000000e+0, 1.000000e+0]> : tensor<2xf32>
  %1 = stablehlo.transpose %0, dims = [0] : (tensor<2xf32>) -> tensor<2xf32>
  return %1 : tensor<2xf32>
}
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant dense<[0.000000e+00, 1.000000e+00]> : tensor<2xf32>
// CHECK-NOT: transpose
// CHECK: return %[[CONST_0]] : tensor<2xf32>

// -----

  %1 = stablehlo.constant dense<[4.0, 6.0]> : tensor<2xf32>
  %2 = stablehlo.divide %0, %1 : tensor<2xf32>
  return %2 : tensor<2xf32>
}

// CHECK-LABEL: foldDivLHSSplat
// CHECK: stablehlo.constant dense<[5.000000e-01, 0.333333343]> : tensor<2xf32>

// -----

func.func @foldDivRHSSplat() -> tensor<2xf32> {
  %0 = stablehlo.constant dense<[4.0, 6.0]> : tensor<2xf32>
  %1 = stablehlo.constant dense<2.0> : tensor<2xf32>

  func.return %1 : tensor<4x4x4xf32>
}

func.func @testMultipleLstms(%arg0: tensor<4x4xf32>, %arg1: tensor<4xf32>, %arg2: tensor<4x4x4xf32>) -> tensor<4x4x4xf32> {

  // CHECK: %{{.*}} = "tfl.reshape"(%{{.*}}, %[[cst]]) : (tensor<4xf32>, tensor<2xi32>) -> tensor<2x2xf32>
  %cst = arith.constant dense<[2, 2]> : tensor<2xi32>
  %0 = "tfl.pseudo_const" () {value = dense<1.0> : tensor<4xf32>} : () -> tensor<4xf32> loc("Const")
  %1 = "tfl.reshape" (%0, %cst) : (tensor<4xf32>, tensor<2xi32>) -> tensor<2x2xf32> loc("reshape")
  func.return %1 : tensor<2x2xf32>

    "tfl.yield"(%1#0, %1#1) : (tensor<*xi32>, tensor<*xf32>) -> ()
  }) : (tensor<i32>, tensor<1xf32>) -> (tensor<i32>, tensor<1xf32>) loc("WhileOp")
  // CHECK: (tensor<i32>, tensor<1xf32>, tensor<i32>) ->
  // CHECK-SAME: (tensor<i32>, tensor<1xf32>, tensor<i32>)
  func.return %0#1 : tensor<1xf32>
}

func.func private @WhileOp_cond(%arg0: tensor<*xi32>, %arg1: tensor<*xf32>, %arg2: tensor<i32>) -> tensor<i1> {

    %3 = "tf.XlaCallModule"(%2, %1, %0) <{Sout = [#tf_type.shape<?x2>], module = "", version = 9 : i64}> {_entry_function = @composite_dot_general_with_bias_and_relu6_dynamic_fn_1, _original_entry_function = "composite_dot_general_with_bias_and_relu6_dynamic_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<?x2xf32>, tensor<2x2xf32>, tensor<2xf32>) -> tensor<?x2xf32>

      %cst_1 = arith.constant dense<2.0> : tensor<3xf32>
      %0 = "tf.XlaCallModule"(%arg0, %cst_0, %cst_1) <{Sout = [#tf_type.shape<1x4x4x3>], module = "", version = 9 : i64}> {_entry_function = @composite_fn_1, _original_entry_function = "composite_fn_1", _tfl_quant_trait = "fully_quantizable"} : (tensor<1x4x4x3xf32>, tensor<3x1x1x3xf32>, tensor<3xf32>) -> tensor<1x4x4x3xf32>

  %biasadd = "tf.BiasAdd"(%conv, %dq_bias) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
  %res = "tf.Relu6"(%biasadd) : (tensor<*xf32>) -> tensor<*xf32>
  %q_res = "quantfork.qcast"(%res) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.023529411764705882:-128>>

    module {
      func.func @test_fold_constant(%arg0: tensor<*xf32>) -> (tensor<?x?x?x3xf32>) {
        %cst = "tf.Const"() {value = dense<2.000000e+00> : tensor<2x3x3x1xf32>} : () -> tensor<2x3x3x1xf32>
        %cst_0 = "tf.Const"() {value = dense<0.400000e+00> : tensor<3xf32>} : () -> tensor<3xf32>
        %cst_1 = "tf.Const"() {value = dense<0.500000e+00> : tensor<3xf32>} : () -> tensor<3xf32>

  %2 = func.call @fold_all_test(%arg0, %0, %1) : (tensor<256x32x32x3xf32>, tensor<16x3x3x3xf32>, tensor<16xf32>) -> tensor<256x30x30x16xf32>
  func.return %2 : tensor<256x30x30x16xf32>
}

// ALL-LABEL: @fold_all_test
func.func @fold_all_test(%arg0: tensor<256x32x32x3xf32>, %arg1: tensor<16x3x3x3xf32>, %arg2: tensor<16xf32>) -> tensor<256x30x30x16xf32> {