// CHECK:           %[[VAL_4:.*]] = "tf.Transpose"(%[[VAL_0]], %[[VAL_3]]) : (tensor<2x3xf32>, tensor<2xi64>) -> tensor<3x2xf32>
// CHECK:           return %[[VAL_4]] : tensor<3x2xf32>
// CHECK:         }
func.func @transpose_2d(%arg0: tensor<2x3xf32>) -> tensor<3x2xf32> {
  %0 = mhlo.constant dense<[1, 0]> : tensor<2xi64>
  %1 = mhlo.constant dense<[1, 0]> : tensor<2xi64>

  %1 = "tfl.transpose"(%arg0, %perm) : (tensor<1x40x40x3xf32>, tensor<4xi32>) -> tensor<1x3x40x40xf32>
  %2 = "tfl.add"(%1, %bias) {fused_activation_function = "NONE"} : (tensor<1x3x40x40xf32>, tensor<1x3x1x1xf32>) -> tensor<1x3x40x40xf32>
  func.return %2 : tensor<1x3x40x40xf32>

  // CHECK: %[[transpose:.*]] = "tfl.transpose"
  // CHECK: %[[add:.*]] = tfl.add(%[[transpose]],
  // CHECK: return %[[add]]
}

  %0 = "tf.Transpose"(%arg0, %permutation) : (tensor<1x2x3xf32>, tensor<3xi32>) -> tensor<3x2x1xf32>
  func.return %0 : tensor<3x2x1xf32>
}

// -----

// CHECK-LABEL: @transpose_3d
func.func @transpose_3d(%arg0: tensor<1x2x3xf32>) -> tensor<3x2x1xf32> {
  %permutation = "tf.Const"() {value = dense<[2, 1, 0]> : tensor<3xi64>} : () -> (tensor<3xi64>)
  // CHECK: "mhlo.transpose"

//===----------------------------------------------------------------------===//

// If the 'transpose' attribute is true returns ElementsAttr to transpose 2D
// matrix. Otherwise, returns ElementsAttr for identity transpose.
static DenseIntElementsAttr Get2DTransposePerm(BoolAttr transpose, Builder *b) {
  if (transpose.getValue()) return GetI64ElementsAttr({1, 0}, b);
  return GetI64ElementsAttr({0, 1}, b);
}

  %86 = "tf.Range"(%12, %85, %15) {device = ""} : (tensor<i64>, tensor<i64>, tensor<i64>) -> tensor<?xi64>
  %87 = "tf.Reshape"(%86, %4) {device = ""} : (tensor<?xi64>, tensor<2xi32>) -> tensor<2x?xi64>
  %88 = "tf.Transpose"(%87, %8) {device = ""} : (tensor<2x?xi64>, tensor<2xi32>) -> tensor<?x2xi64>
  %89 = "tf.Reshape"(%88, %6) {device = ""} : (tensor<?x2xi64>, tensor<1xi32>) -> tensor<?xi64>