func.return %0 : tensor<8x16xf32>
}

// -----

func.func @testCumprod(%arg: tensor<8x16xf32>) -> tensor<8x16xf32> {
  %axis = arith.constant dense<-3> : tensor<i32>
  // expected-error @+1 {{axis operand should be within range [-2, 2)}}
  %0 = "tf.Cumprod"(%arg, %axis) : (tensor<8x16xf32>, tensor<i32>) -> tensor<8x16xf32>
  func.return %0 : tensor<8x16xf32>
}

    %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} : () -> tensor<10xf32>
    %cst_0 = "tf.Const"() {value = dense<[-1, 10]> : tensor<2xi32>} : () -> tensor<2xi32>
    %1 = "tf.MatMul"(%arg0, %arg1) {
      transpose_a = false, transpose_b = false

  %one = "glue.constant"() { value = 1: i32 } : () -> i32
  %done = "glue.compare" (%one, %one) { predicate = #glue<"compare LTE"> } : (i32, i32) -> i1
  %2 = mhlo.constant dense<[[1.1]]> : tensor<1x1xf32>
  %3 = mhlo.multiply %2, %2 : tensor<1x1xf32>
  %cst = "tf.Const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
  %0 = "tf.AddV2"(%arg0, %cst) {device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>

  constexpr absl::string_view kProgram = R"mlir(
func.func @main(
    %arg0: tensor<10x10xf32>, %scale: tensor<10xf32>, %zp: tensor<10xi32>
  ) -> tensor<10x10xi8> {
  %0 = "tf.UniformQuantize"(%arg0, %scale, %zp) {
    quantization_axis = 1 : i64,
    quantization_min_val = -128 : i64,
    quantization_max_val = 127 : i64
  } : (tensor<10x10xf32>, tensor<10xf32>, tensor<10xi32>) -> tensor<10x10x!tf_type.qint8>

func.func @QuantizedCatsAddRequantsTest(%arg0: tensor<1x1xf32>, %arg1: tensor<1x1xf32>, %arg2: tensor<1x1xf32>, %arg3: tensor<1x1xf32>) -> (tensor<1x4xf32>, tensor<1x3xf32>) {
  %0 = "quantfork.stats"(%arg0) {layerStats = dense<[-0.440728068, 0.189515018]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>
  %1 = "quantfork.stats"(%arg1) {layerStats = dense<[-0.154693216, 0.26483655]> : tensor<2xf32>} : (tensor<1x1xf32>) -> tensor<1x1xf32>

// CHECK: return %[[identity]] : tensor<1x3x2x2xf32>

func.func @cast_bf16_matmul_to_fp32(%arg0: tensor<1x10xf32>) -> (tensor<1x2xf32>) {
  %cst = "tf.Const"() {device = "", value = dense<1.000000e+01> : tensor<10x2xbf16>} : () -> tensor<10x2xbf16>
  %0 = "tf.Cast"(%arg0) {Truncate = false, device = ""} : (tensor<1x10xf32>) -> tensor<1x10xbf16>

        unknown_rank: true
      }
    }
  }
  attr {
    key: "dtype"
    value {
      type: DT_BOOL
    }
  }
}

# CHECK-DAG: tf.IteratorGetNext{{.+}}-> (tensor<1x8xbf16>, tensor<2x?x16xf32>, tensor<*xf64>)
# CHECK-DAG: tf.IteratorGetNextSync{{.+}}-> (tensor<*xi16>, tensor<3x24xi32>, tensor<?x4x32xi64>)

    "tfl.yield"(%1#0, %1#1, %1#2) : (tensor<1x?x16xf32>, tensor<1x!tf_type.string>, tensor<1xi32>) -> ()
  }) {_tfl_quant_trait = "fully_quantizable"} : (tensor<1x!tf_type.string>) -> (tensor<1x?x16xf32>, tensor<1x!tf_type.string>, tensor<1xi32>)

    ```mlir
      %0 = "tf.Const"() {value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
      %1 = "tf.Const"() {device = "", value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
      %2 = "tf.Const"() {device = "baz", value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
    ```

    then running this pass with 'default-device=foobar', we get:

    ```mlir

  %0 = "mhlo.reshape"(%arg0) : (tensor<1x1x512xf32>) -> tensor<1x512xf32>
  %1 = "mhlo.dot"(%0, %arg1) : (tensor<1x512xf32>, tensor<512x13x!quant.uniform<i8:f32, 0.00285>>) -> tensor<1x13xf32>
  %2 = "mhlo.reshape"(%1) : (tensor<1x13xf32>) -> tensor<1x1x13xf32>
  func.return %2 : tensor<1x1x13xf32>

// CHECK:      %[[RES:.*]] = "mhlo.dot_general"(%arg0, %arg1) <{
// CHECK-SAME:   dot_dimension_numbers = #mhlo.dot<