// Identity function to make the exporter happy
func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> {
  func.return %arg0 : tensor<4xi8>
}

//CHECK:func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> attributes {tf.entry_function = {inputs = "arg0", outputs = "arg0"}} {
//CHECK: return %arg0 : tensor<4xi8>
//CHECK:}

// CHECK: %[[maxpool_i8:.*]] = "tf.MaxPool"(%[[sc1]])
// CHECK-SAME: (tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xi8>, tensor<2xi32>) -> tensor<*xi8>
// CHECK: %[[sc2:.*]] = "quantfork.scast"(%[[reshape_i8]])
// CHECK: %[[dq:.*]] = "quantfork.dcast"(%[[sc2]]) : (tensor<*x!quant.uniform<i8:f32, 5.000000e-02:-10>>)

  func.return %11 : tensor<*xf32>
}

// CHECK: %[[maxpool_i8:.*]] = "tf.MaxPool"
// CHECK-SAME: (tensor<*xi8>) -> tensor<*xi8>
// CHECK: %[[reshape_i8:.*]] = "tf.Reshape"(%[[maxpool_i8]]
// CHECK-SAME: (tensor<*xi8>, tensor<2xi32>) -> tensor<*xi8>
// CHECK: %[[scast:.*]] = "quantfork.scast"(%[[reshape_i8]]
// CHECK: %[[matmul:.*]] = "tf.PartitionedCall"(%[[scast]]
// CHECK-SAME: f = @composite_matmul_fn_1

  // Argument has rank 2, but its dimensions are unknown.
  constexpr absl::string_view kArgumentWithUnknownDims = R"mlir(
    func.func @unknown_dims_arg(%arg: tensor<?x?xi8>) -> tensor<?x?xi8> {
      return %arg : tensor<?x?xi8>
    }
  )mlir";

  OwningOpRef<ModuleOp> module_op =
      ParseModuleOpString(kArgumentWithUnknownDims);
  ASSERT_TRUE(module_op);

  %7 = "tf.opH"() : () -> tensor<*xf32>
  %8 = "tf.opI"() : () -> tensor<*xf32>
  %9 = "tf.opJ"() : () -> tensor<*xi8>
  %10 = "tf.opK"() : () -> tensor<*xi16>
  %11 = "tf.opL"() : () -> tensor<*xi64>

// test stablehlo roundtrip

// Identity function to make the exporter happy
func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> {
  func.return %arg0 : tensor<4xi8>
}

//CHECK:func.func @main(%arg0: tensor<4xi8>) -> tensor<4xi8> attributes {tf.entry_function = {inputs = "arg0", outputs = "arg0"}} {
//CHECK: return %arg0 : tensor<4xi8>
//CHECK:}

func.func @logistic(%arg0: tensor<1x1x1x96xf32>) -> tensor<1x1x1x96xf32> {

  %f0 = "tfl.logical_and"(%arg0, %cst_false0): (tensor<1xi1>, tensor<1xi1>) -> tensor<1xi1>        // match
  %f1 = "tfl.logical_and"(%arg0, %cst_false1): (tensor<1xi1>, tensor<2x3xi1>) -> tensor<2x3xi1>    // match
  %f2 = "tfl.logical_and"(%arg1, %cst_false0): (tensor<2x3xi1>, tensor<1xi1>) -> tensor<2x3xi1>    // no match

}

TEST_F(ConvertTfQuantToMhloIntTest, UniformRequantizeToValidGraph) {
  constexpr absl::string_view kProgram = R"mlir(
func.func @main(%input: tensor<10xi8>) -> tensor<10xi8> {
  %input_scale = "tf.Const"() { value = dense<0.2235> : tensor<f32> } : () -> tensor<f32>
  %input_zp = "tf.Const"() { value = dense<-2> : tensor<i32> } : () -> tensor<i32>

          )
      )
      module_str = self._extract_first_xla_call_module_op(
          self._output_saved_model_path
      )
      self.assertTrue(
          re.search('stablehlo.dot_general.*xi8>.*xi8>.*xi32>', module_str)
      )

    # Run quantization the first time, calibration is expected to be skipped.
    output_saved_model_path_2 = self.create_tempdir('output2').full_path

  if (!quantDenseAttr) {
    return nullptr;
  }

  // Cast from an expressed-type-based type to storage-type-based type,
  // preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>).
  ShapedType newSparseType = mlir::dyn_cast_or_null<ShapedType>(
      quantizedElementType.castExpressedToStorageType(
          realSparseAttr.getType()));
  if (!newSparseType) {
    return nullptr;
  }