/*error_tolerance=*/1.0);
}

TEST_F(ConvertTfQuantToMhloIntTest, UniformQuantizeAddToValidGraph) {
  constexpr absl::string_view kProgram = R"mlir(
func.func @main(%lhs: tensor<10x10xi32>, %rhs: tensor<10x10xi32>) -> tensor<10x10xi32> {
  %lhs_scale = "tf.Const"() { value = dense<0.518> : tensor<f32> } : () -> tensor<f32>
  %lhs_zp = "tf.Const"() { value = dense<42> : tensor<i32> } : () -> tensor<i32>

  func.func @random_uniform_int(%arg0: tensor<i32>, %arg1: tensor<i32>) -> tensor<1000xi32> {
    %0 = "tf.Const"() {value = dense<1000> : tensor<1xi32>} : () -> tensor<1xi32>
    // CHECK-NOT: tf.RandomUniformInt
    %1 = "tf.RandomUniformInt"(%0, %arg0, %arg1) {seed = 0 : i64, seed2 = 0 : i64} : (tensor<1xi32>, tensor<i32>, tensor<i32>) -> tensor<1000xi32>
    func.return %1 : tensor<1000xi32>
  }

  // CHECK-LABEL: multinomial

  %output_state, %output = "vhlo.rng_bit_generator_v1"(%arg0) <{rng_algorithm = #vhlo<rng_algorithm_v1 DEFAULT>}> : (tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>)
  func.return %output_state, %output : tensor<2xui64>, tensor<10x12xui32>
}

// CHECK:func.func private @rng_bit_generator(%arg0: tensor<2xui64>) -> (tensor<2xui64>, tensor<10x12xui32>) {

// CHECK-LABEL: float_matmul
func.func @float_matmul(
  %arg0: tensor<1x10xf32>, %arg1: tensor<10x10xf32>) -> (tensor<*xf32>, tensor<*xf32>, tensor<*xf32>) {
  %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} : () -> tensor<10xf32>
  %0 = "tf.MatMul"(%arg0, %arg1) {
    transpose_a = false, transpose_b = false
  } : (tensor<1x10xf32>, tensor<10x10xf32>) -> tensor<*xf32>

  "tfl.custom"(%0, %2) {custom_code = "FlexAssignVariableOp", custom_option = #tfl<const_bytes : "0x1041737369676E5661726961626C654F70003B121041737369676E5661726961626C654F701A001A002A0B0A056474797065120230012A140A0E76616C69646174655F736861706512022800320000024F3E1414042801">} : (tensor<!tf_type.resource<tensor<1x10xf32>>>, tensor<1x10xf32>) -> ()

  // CHECK-NEXT: %[[UPDATE:.*]] = "tf.XlaDynamicUpdateSlice"(%[[STACK_VAL]], %[[UPDATE_SLICE]], %[[CONCAT_OFFETS]]) : (tensor<10x2xi32>, tensor<1x2xi32>, tensor<2xi32>) -> tensor<10x2xi32>
  // CHECK-NEXT: "tf.AssignVariableOp"(%[[BUFFER]], %[[UPDATE]]) : (tensor<!tf_type.resource<tensor<10x2xi32>>>, tensor<10x2xi32>) -> ()
  // CHECK-NEXT: %[[CONST1:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}> : () -> tensor<1xi32>

// -----

func.func @constant_tensor_invalid_3(%arg0: vector<1xi32>) -> tensor<1x1xi32> {
    // expected-error@+1 {{input and output should have same shape and element type}}
  %0 = "tfr.constant_tensor"(%arg0) : (vector<1xi32>) -> tensor<1x1xi32>
  func.return %0 : tensor<1x1xi32>
}

// -----

func.func @constant_tensor_invalid_4(%arg0: i32) -> tensor<1x1xi32> {
    // expected-error@+1 {{input can not be converted to an output tensor}}

// CHECK-DAG: %[[SCALE:.*]] = "tf.Const"() <{value = dense<0.0787401571> : tensor<f32>
}

// -----

module {
  func.func @quantize_xlagather(%arg0: tensor<10x2xi32>) -> tensor<1x300x10xf32> {
    %cst_0 = "tf.Const"() {device = "", value = dense<1.000000e+01> : tensor<200x100x300xf32>} : () -> tensor<200x100x300xf32>

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

  %const_64_1 = "tf.Const"() {value = dense<[64, 1]> : tensor<2xi32>} : () -> tensor<2xi32>
  %1 = "tf.Range"(%const_0, %const_10000, %const_1) : (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<10000xi32>
  %2 = "tf.ExpandDims"(%1, %const_0) : (tensor<10000xi32>, tensor<i32>) -> tensor<1x10000xi32>
  %3 = "tf.Tile"(%2, %const_64_1) : (tensor<1x10000xi32>, tensor<2xi32>) -> tensor<64x10000xi32>
  // CHECK-NOT: tf.Range
  // CHECK-NOT: tf.ExpandDims