%cst = arith.constant dense<[[1], [2]]> : tensor<2x1xi32>

  // CHECK: %[[CST:.*]] = arith.constant dense<2> : tensor<1xi32>
  // CHECK: return %[[CST]]
  %0 = "tfl.rank"(%cst) : (tensor<2x1xi32>) -> tensor<1xi32>
  func.return %0 : tensor<1xi32>
}

// CHECK-LABEL: @rank_input_known_rank
func.func @rank_input_known_rank(%arg0 : tensor<2x1xi32>) -> tensor<1xi32> {

  // CHECK-NEXT: %12 = "mhlo.pad"(%10, %11) <{edge_padding_high = dense<[3, 4]> : tensor<2xi64>, edge_padding_low = dense<0> : tensor<2xi64>, interior_padding = dense<0> : tensor<2xi64>}> : (tensor<3x4xi32>, tensor<i32>) -> tensor<6x8xi32>
  // CHECK-NEXT: %13 = "mhlo.slice"(%arg1) <{limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}> : (tensor<2xi32>) -> tensor<1xi32>

}

func.func @concat_v2_i64_axis(%arg0: tensor<2x1xi32>, %arg1: tensor<2x1xi32>, %arg2: tensor<2x1xi32>) -> tensor<2x3xi32> {
  %0 = "tf.Const"() { value = dense<-1> : tensor<i64> } : () -> tensor<i64>
  %1 = "tf.ConcatV2"(%arg0, %arg1, %arg2, %0) : (tensor<2x1xi32>, tensor<2x1xi32>, tensor<2x1xi32>, tensor<i64>) -> tensor<2x3xi32>

  %19 = "tf.StringLength"(%arg0) {device = "", unit = "BYTE"} : (tensor<1x!tf_type.string>) -> tensor<1xi32>
  %20 = "tf.ExpandDims"(%19, %7) {device = ""} : (tensor<1xi32>, tensor<i32>) -> tensor<1x1xi32>
  %21 = "tf.Cast"(%20) {Truncate = false, device = ""} : (tensor<1x1xi32>) -> tensor<1x1xi64>
  %22 = "tf.Reshape"(%21, %12) {device = ""} : (tensor<1x1xi64>, tensor<1xi64>) -> tensor<1xi64>

// CHECK-DAG: %[[CONST_0:.+]] = "tfl.pseudo_const"() <{value = dense<{{\[\[0, 0\], \[1, 1\], \[1, 1\], \[0, 0\]\]}}> : tensor<4x2xi32>}> : () -> tensor<4x2xi32>

  %cst_2 = arith.constant dense<0> : tensor<4x2xi32>
  %0 = "tf.SpaceToBatchND"(%arg0, %cst, %cst_0) : (tensor<1x128x128x3xf32>, tensor<2xi32>, tensor<2x2xi32>) -> tensor<4x68x68x3xf32>
  %1 = "tf.Conv2D"(%0, %arg1) {padding = "VALID", strides = [1, 1, 1, 1]} : (tensor<4x68x68x3xf32>, tensor<5x5x3x8xf32>) -> tensor<4x64x64x8xf32>
  %2 = "tf.Pad"(%1, %cst_2) : (tensor<4x64x64x8xf32>, tensor<4x2xi32>) -> tensor<4x64x64x8xf32>

    func.return
  }

  // CHECK: private @__inference_Dataset_map_lambda_96150(%arg0: tensor<64xi32>
  // CHECK-SAME: %arg1: tensor<69xi32>, %arg2: tensor<73xi32>, %arg3: tensor<12xi64>
  // CHECK-SAME: %arg8: tensor<64xi32>, %arg9: tensor<64xi32>
  // CHECK-SAME: %arg12: tensor<10xi32>)
  // CHECK-SAME: -> (tensor<64xi32>, tensor<69xi32>, tensor<10xi32>)

  %11 = "tf.Reshape"(%10, %0) {device = ""} : (tensor<4x1x1x2x1xi32>, tensor<2xi32>) -> tensor<4x2xi32>
  %12 = "tf.Transpose"(%11, %6) {device = ""} : (tensor<4x2xi32>, tensor<2xi32>) -> tensor<2x4xi32>
  %13 = "tf.Reshape"(%arg0, %1) {device = ""} : (tensor<1x1x2x1xf32>, tensor<1xi32>) -> tensor<2xf32>

  // CHECK: %[[IND_RESHPE:.*]] = "tf.Reshape"(%[[ARG1]], %[[IND_SHAPE]]) : (tensor<5xi32>, tensor<2xi32>) -> tensor<5x1xi32>
  // CHECK: %[[SC:.*]] = "tf.TensorScatterUpdate"(%[[BUFFER]], %[[IND_RESHPE]], %[[ARG2]]) : (tensor<10x8x9xf32>, tensor<5x1xi32>, tensor<5x8x9xf32>) -> tensor<10x8x9xf32>

}

// -----

func.func @testPadUnsupportedPaddings(tensor<*xf32>, tensor<6x3xi32>) -> tensor<? x f32> {
^bb0(%arg0: tensor<*xf32>, %arg1: tensor<6x3xi32>):
  // expected-error @+1 {{'tfl.pad' op failed to verify that the first dim size of the padding argument must be at most 5}}
  %0 = "tfl.pad"(%arg0, %arg1) : (tensor<*xf32>, tensor<6x3xi32>) -> tensor<? x f32>
  func.return %0#0 : tensor<? x f32>
}

// -----