}

// -----
// This test verifies that the pipelining works correctly for evaluation-only models.
module {
  func.func @main() {
    %cst = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
    %0 = "tf.While"(%cst) {body = @while_body, cond = @while_cond, is_stateless = false} : (tensor<i32>) -> (tensor<i32>)
    return
  }
  func.func private @while_body(%arg0: tensor<i32>) -> (tensor<i32>) {

    %cst = arith.constant dense<1> : tensor<4xi32>
    %cst_0 = arith.constant dense<2> : tensor<1xi32>
    %cst_1 = arith.constant dense<[2, 1]> : tensor<2xi32>
    %0 = "tfl.reshape"(%arg0, %cst) {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1xf32>, tensor<4xi32>) -> tensor<1x1x1x1xf32>

  %5 = "quantfork.stats"(%4) {layerStats = dense<[-1.0, 1.0]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>
  func.return %5 : tensor<1x1xf32>

// CHECK: %[[cst:.*]] = arith.constant dense<{{.*}}> : tensor<1xf32>
// CHECK-NEXT: %[[q1:.*]] = "tfl.quantize"(%[[cst]]) <{qtype = tensor<1x!quant.uniform<i32:f32, {{.*}}>>}> {volatile} : (tensor<1xf32>) -> tensor<1x!quant.uniform<i32:f32, {{.*}}>>

      func.func @main() -> (tensor<32x64xf32> {mhlo.sharding = "\08\01\1A\01\01\22\01\00"}) {
        %cst = "tf.Const"() {value = dense<[524170, 523952]> : tensor<2xi32>} : () -> tensor<2xi32>
        %cst_0 = "tf.Const"() {value = dense<[32, 64]> : tensor<2xi32>} : () -> tensor<2xi32>
        %0 = "tf.StatelessRandomNormal"(%cst_0, %cst) : (tensor<2xi32>, tensor<2xi32>) -> tensor<32x64xf32>
        return %0 : tensor<32x64xf32>
    }
  })";

  // CHECK:      })
  // CHECK-SAME: (tensor<3x3xf32>, tensor<3x2xi32>, tensor<3xf32>) -> tensor<3x3xf32>

  // return %[[RESULT]] : tensor<3x3xf32>

    %cst = mhlo.constant dense<3> : tensor<2xi32>
    %0 = "tf.SparseToDense"(%arg0, %cst, %arg1, %arg2) {validate_indices = true}: (tensor<3x2xi32>, tensor<2xi32>, tensor<3xf32>, tensor<f32>) -> tensor<3x3xf32>
    func.return %0 : tensor<3x3xf32>
  }

      %cst = "tf.Const"() {value = dense<[1, -1, 4]> : tensor<3xi32>} : () -> tensor<3xi32>

// CHECK-LABEL: testWhileRegionExternInCond
func.func @testWhileRegionExternInCond(%arg0 : tensor<*xf32>, %arg1 : tensor<i32>, %arg2 : tensor<i32>) -> tensor<*xf32> {
  %cst = arith.constant dense<4> : tensor<i32>
  %limit = "tf.Add"(%arg2, %cst) : (tensor<i32>, tensor<i32>) -> tensor<i32>
  // CHECK: [[Result:%.*]]:3 = "tf.While"(%arg0, %arg1, %{{.+}} <{body = @tf.WhileRegion_body, cond = @tf.WhileRegion_cond

// CHECK:         }

// -----

func.func @ensureBiasForConv2d(%arg0: tensor<128x32x32x3xf32>, %arg1: tensor<32x1x1x3xf32>) -> tensor<128x32x32x32xf32> {
  %cst = "tfl.no_value"() {value = unit} : () -> none

  func.return %0 : tensor<1x1xi32>
}

// -----

// CHECK-LABEL: get_element
func.func @get_element(%arg0: !tfr.tensor_list<T>) -> !tfr.tensor {
  %cst = "arith.constant"() {value = 1 : index} : () -> index
  %0 = tfr.get_element %arg0[%cst] : (!tfr.tensor_list<T>, index) -> !tfr.tensor
  func.return %0 : !tfr.tensor
}

// -----

// CHECK-LABEL: build_list

    module attributes {tf.versions = {bad_consumers = [], min_consumer = 0 : i32, producer = 268 : i32}} {
      func.func @main() -> (tensor<2xi32>) {
        %cst = "tf.Const"() {value = dense<[524170, 523952]> : tensor<2xi32>} : () -> tensor<2xi32>
        return %cst : tensor<2xi32>
    }
  })";

  auto compilation_result = CompileMlirModule(
      kHasReturnValuesAndNoMetadataRetvals,