// 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

// CHECK-DAG: %[[CONST:.*]] = "tf.Const"() <{value = dense<-131072> : tensor<1x3xi32>}> : () -> tensor<1x3xi32>
// CHECK: %[[MATMUL:.*]] = "tf.XlaDotV2"({{.*}}, %[[WEIGHT]])
// CHECK-SAME: (tensor<1x1024xi8>, tensor<1024x3xi8>) -> tensor<1x3xi32>
// CHECK: %[[SUB:.*]] = "tf.Sub"(%[[MATMUL]], %[[CONST]]) : (tensor<1x3xi32>, tensor<1x3xi32>) -> tensor<1x3xi32>
}

// -----

module attributes {} {

  %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>

}

// CHECK-LABEL: testRedundantReshape
func.func @testRedundantReshape(%arg0: tensor<4x4xi32>) -> tensor<2x8xi32> {
  %0 = "tf.Const"() {value = dense<[8, 2]> : tensor<2xi32>} : () -> tensor<2xi32>
  %1 = "tf.Const"() {value = dense<[2, 8]> : tensor<2xi32>} : () -> tensor<2xi32>
  %2 = "tf.Reshape"(%arg0, %0) : (tensor<4x4xi32>, tensor<2xi32>) -> tensor<8x2xi32>

func.func @scatter_nd_add(%arg0: tensor<7xi64>, %arg1: tensor<1x1xi32>, %arg2: tensor<1xi64>) -> tensor<7xi64> {
  %0 = "tf.TensorScatterAdd"(%arg0, %arg1, %arg2) : (tensor<7xi64>, tensor<1x1xi32>, tensor<1xi64>) -> tensor<7xi64>
  func.return %0 : tensor<7xi64>

  // CHECK-LABEL: scatter_nd_add
  // CHECK:  %[[GATHER:.*]] = "tf.GatherNd"(%arg0, %arg1) <{bad_indices_policy = ""}> : (tensor<7xi64>, tensor<1x1xi32>) -> tensor<1xi64>

  // CHECK-NEXT: %[[UPDATE_SLICE:.*]] = "tf.Reshape"(%[[PUSH_VAL]], %[[UPDATE_SHAPE]]) : (tensor<2xi32>, tensor<2xi32>) -> tensor<1x2xi32>
  // CHECK-NEXT: %[[ZERO_INDS:.*]] = "tf.Const"() <{value = dense<0> : tensor<1xi32>}> : () -> tensor<1xi32>
  // CHECK-NEXT: %[[CONCAT_DIM:.*]] = "tf.Const"() <{value = dense<0> : tensor<i32>}> : () -> tensor<i32>

// CHECK: %[[CONST_0:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<2x3xi8>}> : () -> tensor<2x3x!quant.uniform<i8<-127:127>:f32:1, {{.*}}>
// CHECK: %[[CONST_1:.+]] = stablehlo.constant() <{value = dense<{{.*}}> : tensor<1x3xi32>}> : () -> tensor<1x3x!quant.uniform<i32:f32:1, {{.*}}>
// CHECK: %[[UNIFORM_QUANTIZE_0:.+]] = stablehlo.uniform_quantize %[[ARG_0]] : (tensor<1x2xf32>) -> tensor<1x2x!quant.uniform<i8:f32, {{.*}}>>

// -----

// Tests that a float `stablehlo.reshape` is not converted to `tfl.reshape`.

func.func @reshape_float(%arg0: tensor<2x3x4xf32>) -> tensor<6x4xf32> {
  %0 = stablehlo.reshape %arg0 : (tensor<2x3x4xf32>) -> tensor<6x4xf32>
  return %0 : tensor<6x4xf32>
}
// CHECK-LABEL: reshape_float
// CHECK-NOT: tfl.reshape
// CHECK: stablehlo.reshape

// -----