// CHECK:}

func.func @reshape(%arg0 : tensor<1x128xi32>) -> tensor<4x32x1xi32>{
  %0 = "vhlo.reshape_v1"(%arg0) : (tensor<1x128xi32>) -> tensor<4x32x1xi32>
  func.return %0 : tensor<4x32x1xi32>
}

//CHECK:func.func private @reshape(%arg0: tensor<1x128xi32>) -> tensor<4x32x1xi32> {
//CHECK-NEXT: %0 = "vhlo.reshape_v1"(%arg0) : (tensor<1x128xi32>) -> tensor<4x32x1xi32>
//CHECK-NEXT: return %0 : tensor<4x32x1xi32>

  func.return %4 : tensor<8x128xf32>

// CHECK-LABEL: SoftMaxWithNormalization
// CHECK: %[[RESULT:.*]] = "tfl.softmax"(%arg0) <{beta = 1.000000e+00 : f32}> : (tensor<8x128xf32>) -> tensor<8x128xf32>
// CHECK: return %[[RESULT]] : tensor<8x128xf32>
}

func.func @SoftMaxWithoutNormalization(%arg0: tensor<8x128xf32>) -> tensor<8x128xf32> {

    %11 = stablehlo.convert %2 : (tensor<1x1x1xi8>) -> tensor<1x1x1xi32>
    %12 = stablehlo.broadcast_in_dim %11, dims = [0, 1, 2] : (tensor<1x1x1xi32>) -> tensor<8x16x16xi32>
    %13 = stablehlo.subtract %10, %12 : tensor<8x16x16xi32>  // q1 - z1
    %14 = stablehlo.convert %9 : (tensor<8x16x4xi8>) -> tensor<8x16x4xi32>
    %15 = stablehlo.convert %4 : (tensor<1x1x1xi8>) -> tensor<1x1x1xi32>

  func.return %0 : tensor<i32>
}

// -----

// CHECK-LABEL: constant_tensor_vector
func.func @constant_tensor_vector(%arg0: vector<1x2xi32>) -> tensor<1x2xi32> {
  %0 = "tfr.constant_tensor"(%arg0) : (vector<1x2xi32>) -> tensor<1x2xi32>
  func.return %0 : tensor<1x2xi32>
}

// -----

// CHECK-LABEL: constant_tensor_array
func.func @constant_tensor_array() -> !tfr.tensor {

// CHECK:           %[[SUB:.*]] = "tf.Sub"(%[[CST]], %[[SCATTER]]) : (tensor<i32>, tensor<1x24xi32>) -> tensor<1x24xi32>
// CHECK:           %[[MUL:.*]] = "tf.Mul"(%[[SUB]], %[[CAST0]]) : (tensor<1x24xi32>, tensor<1x24xi32>) -> tensor<1x24xi32>

  %0 = "mhlo.broadcast_in_dim"(%arg1) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}> : (tensor<1x1xi32>) -> tensor<1x100xi32>
  %1 = "mhlo.select"(%arg0, %0, %arg2) : (tensor<i1>, tensor<1x100xi32>, tensor<1x100xi32>) -> tensor<1x100xi32>
  func.return %1 : tensor<1x100xi32>
}

// CHECK-LABEL:   func @selectv2_broadcasted_condition(

    %0 = "tf.Const"() {value = dense<1> : tensor<i64>} : () -> tensor<i64>
    %1 = "tf.Const"() {value = dense<0> : tensor<1x1xi32>} : () -> tensor<1x1xi32>
    %2 = "tf.Const"() {value = dense<[7, 7, 3, 64]> : tensor<4xi32>} : () -> tensor<4xi32>
    %3 = "tf.Const"() {value = dense<[[0, 0], [3, 3], [3, 3], [0, 0]]> : tensor<4x2xi32>} : () -> tensor<4x2xi32>

module {
    func.func @main(%arg0: tensor<8x128xf32>, %arg1: tensor<f32>) -> tensor<11x131xf32> {
      %0 = "tfl.custom"(%arg0, %arg1) {custom_code = "stablehlo.pad", custom_option = #tfl<const_bytes : "0x656467655F70616464696E675F6869676800020203656467655F70616464696E675F6C6F7700020100696E746572696F725F70616464696E6700020000033E2A17030103311E0B2C2C2C062401">} : (tensor<8x128xf32>, tensor<f32>) -> tensor<11x131xf32>

func.func @testConv2dShapeInferenceDynamic(%arg0: tensor<1x?x?x128xf32>, %arg1: tensor<128x3x3x128xf32>, %arg2: tensor<128xf32>) -> tensor<1x?x?x128xf32> {
  // CHECK: "tfl.conv_2d"(%arg0, %arg1, %arg2) <{dilation_h_factor = 2 : i32, dilation_w_factor = 2 : i32, fused_activation_function = "NONE", padding = "VALID", stride_h = 1 : i32, stride_w = 1 : i32}> : (tensor<1x?x?x128xf32>, tensor<128x3x3x128xf32>, tensor<128xf32>) -> tensor<1x?x?x128xf32>

// CHECK: func @while_cond(%arg0: tensor<256x256xi32>)
// CHECK: func private @while_body_0(%arg0: tensor<128xi32>)
// CHECK: func private @while_cond_1(%arg0: tensor<128xi32>)
func.func @while_main(%arg0: tensor<256x256xi32>, %arg1: tensor<128xi32>) -> (tensor<256x256xi32>, tensor<128xi32>) {