// CHECK-LABEL: reduce_window_max_activation_transpose
// CHECK-SAME: (%[[ARG:.+]]: tensor<16x8xf32>) -> tensor<4x8xf32>
func.func @reduce_window_max_activation_transpose_rank2(%arg0: tensor<16x8xf32>) -> tensor<4x8xf32> {
  %0 = stablehlo.constant dense<0xFF800000> : tensor<f32>  // -inf
  %1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<16x8xf32>) -> tensor<8x16xf32>
  %2 = "stablehlo.reduce_window"(%1, %0) ({

  %3 = "tf.TensorListStack"(%1#3, %cst_0) {device = "/job:localhost/replica:0/task:0/device:CPU:0", num_elements = 2 : i64} : (tensor<!tf_type.variant<tensor<*xf32>>>, tensor<1xi32>) -> tensor<2x8xf32>
  return %3, %2 : tensor<2x8xf32>, tensor<2x8xf32>
}


// -----

// Convert a while with multiple tensor array to map_fn

  %group_assignment = "tf.Const"() { value = dense<[[0],[1]]> : tensor<2x1xi32> } : () -> tensor<2x1xi32>
  // CHECK: "mhlo.all_reduce"
  // CHECK{LITERAL}: replica_groups = dense<[[0], [1]]> : tensor<2x1xi64>
  // CHECK-NOT: channel_handle
  // CHECK: mhlo.add
  %0 = "tf.XlaAllReduce"(%input, %group_assignment) {reduce_op = "Add", mode = "CrossReplica"} : (tensor<f32>, tensor<2x1xi32>) -> tensor<f32>
  func.return %0 : tensor<f32>
}

// -----

  %1 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>

func.func @add_dense_dense_int_mixing_1_n() -> tensor<2x2xi32> {
  %cst_0 = arith.constant dense<[[1, 2]]> : tensor<1x2xi32>
  %cst_1 = arith.constant dense<[[3], [4]]> : tensor<2x1xi32>

  %0 = "tfl.add"(%cst_0, %cst_1) {fused_activation_function = "NONE"} : (tensor<1x2xi32>, tensor<2x1xi32>) -> tensor<2x2xi32>

  func.return %0 : tensor<2x2xi32>

        }) : () -> tensor<f32>
        "tf_device.cluster"() ({
          %group = "tf.Const"() <{value = dense<[[0, 1]]> : tensor<1x2xi32>}> : () -> tensor<1x2xi32>
          %arg1_reduced = "tf.XlaAllReduce"(%arg1_id, %group) <{mode = "CrossReplica", reduce_op = "Add"}> : (tensor<f32>, tensor<1x2xi32>) -> tensor<f32>
          "tf.OpA"(%arg1_reduced) : (tensor<f32>) -> ()
          tf_device.return
        }) : () -> ()

// Note: Argument quantization sequence omitted.
// CHECK: stablehlo.dot_general %{{.+}}, %{{.+}}, contracting_dims = [1] x [0] : (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>

// Note: Result dequantization sequence omitted.
// CHECK: return %{{.+}} : tensor<1x3xf32>

// -----

// Tests that a simple dot_general (lifted as a function) with CustomAggregators

  %0 = "tfl.broadcast_to"(%arg0, %arg1) : (tensor<3xi32>, tensor<2xi32>) -> tensor<3x3xi32>
  return %0 : tensor<3x3xi32>
  // CHECK:  %cst = arith.constant dense<1> : tensor<3x3xi32>
  // CHECK:  %0 = tfl.mul(%arg0, %cst) <{fused_activation_function = "NONE"}> : (tensor<3xi32>, tensor<3x3xi32>) -> tensor<3x3xi32>
  // CHECK:  return %0 : tensor<3x3xi32>
}

  %0 = "tfl.pseudo_qconst"() {qtype = tensor<128x!quant.uniform<i32:f32, 0.7>>, value = dense<0> : tensor<128xi32>} : () -> tensor<128x!quant.uniform<i32:f32, 0.7>>

  // CHECK-DAG: %[[cst:.*]] = arith.constant dense<[7, 6, 5, 4, 3, 2]> : tensor<6xi32>
  // CHECK-DAG: %[[cst_0:.*]] = arith.constant dense<[5, 4, 3, 2, 1, 0]> : tensor<6xi32>
  // CHECK-DAG: %[[cst_1:.*]] = arith.constant dense<[1, 2, 3, 4, 5, 6, 7]> : tensor<7xi32>
  // CHECK: %0 = "tfl.reshape"(%arg0, %[[cst]]) : (tensor<7x6x5x4x3x2x1xf32>, tensor<6xi32>) -> tensor<7x6x5x4x3x2xf32>