%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-NEXT: return %[[OUTPUT_STATE]], %[[OUTPUT]] : tensor<2xui64>, tensor<10x12xui32>
  %output_key, %output = "tf.XlaRngBitGenerator"(%cst_0, %arg0, %cst) : (tensor<i32>, tensor<2xui64>, tensor<2xi32>) -> (tensor<2xui64>, tensor<10x12xui32>)
  func.return %output_key, %output : tensor<2xui64>, tensor<10x12xui32>
}

  %1 = "tf.MatMul"(%arg0, %arg0) {device = "", transpose_a = false, transpose_b = false} : (tensor<10x10xi32>, tensor<10x10xi32>) -> tensor<10x10xi32>
  %2 = "tf.PreventGradient"(%0) : (tensor<10x10xi32>) -> tensor<10x10xi32>
  %3 = "tf.PreventGradient"(%1) : (tensor<10x10xi32>) -> tensor<10x10xi32>
  %4 = "tf.AddV2"(%2, %3) {device = ""} : (tensor<10x10xi32>, tensor<10x10xi32>) -> tensor<10x10xi32>

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

    return %values : tensor<10x10xbf16>
  }

  // CHECK-LABEL: fusedBatchNormV3_noTraining

// CHECK-LABEL:   func @not_convert_remainder_for_uint32(
// CHECK-NOT:       "tf.Mod"
func.func @not_convert_remainder_for_uint32(%arg0: tensor<10x8xui32>, %arg1: tensor<10x8xui32>) -> tensor<10x8xui32> {
  %0 = mhlo.remainder %arg0, %arg1 : tensor<10x8xui32>
  func.return %0 : tensor<10x8xui32>
}

// CHECK-LABEL:   func @not_convert_remainder_for_uint64(
// CHECK-NOT:       "tf.Mod"

    func.return %2 : tensor<1xi32>
  }

  // CHECK-LABEL: func @tensor_cast_partial_infer
  func.func @tensor_cast_partial_infer(%arg0: tensor<?x10xi32>) -> tensor<10x?xi32> {
   // CHECK: return {{.*}} : tensor<10x10xi32>
    %2 = tensor.cast %arg0 : tensor<?x10xi32> to tensor<10x?xi32>
    func.return %2 : tensor<10x?xi32>
  }

  // CHECK-LABEL: operand_pack_unranked