return %3 : tensor<1x2x2x1024xbf16>
  }

  func.func private @composite_dequantize_uniform(%arg0: tensor<*xi8>) -> tensor<*xbf16> {
    %cst = "tf.Const"() {value = dense<1.574710e-02> : tensor<bf16>} : () -> tensor<bf16>
    %0 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<*xi8>) -> tensor<*xbf16>
    %1 = "tf.Mul"(%0, %cst) : (tensor<*xbf16>, tensor<bf16>) -> tensor<*xbf16>
    return %1 : tensor<*xbf16>
  }

tensor<i8>, tensor<f64>, tensor<!tf_type.resource>)

# CHECK-SUBTYPE: func @main(%arg0: tensor<10xi32>, %arg1: tensor<!tf_type.variant<tensor<10xf32>>>, %arg2: tensor<10xi8>, %arg3: tensor<10xf64>, %arg4: tensor<!tf_type.resource<tensor<10xi32>>>) -> (tensor<10xi32>, tensor<!tf_type.variant<tensor<10xf32>>>, tensor<10xi8>, tensor<10xf64>, tensor<!tf_type.resource<tensor<10xi32>>>)...

// 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.func @fixed_element_attribute_invalid(%arg0: tensor<2xf32>) -> tensor<2xi8> {
  %0 = "tfr.cast"(%arg0) : (tensor<2xf32>) -> !tfr.tensor
  // expected-error@+1 {{type i8_ can't be resolved for the signature of the op}}
  %1 = tfr.call @tf__invalid_type_op(%0) : (!tfr.tensor) -> !tfr.tensor
  %2 = "tfr.cast"(%1) : (!tfr.tensor) -> tensor<2xi8>
  func.return %2 : tensor<2xi8>
// CHECK: tfr.call @tf__invalid_type_op

    tf_executor.fetch
  }
  func.return
}

func.func @plain() {
  tf_executor.graph {
    %0:2 = tf_executor.island wraps "tf.Placeholder"() {type = i8} : () -> tensor<16xi8>
    tf_executor.fetch
  }
  func.return

  EXPECT_THAT(num_bytes, Eq(8));
}

TEST_F(GetSizeInBytesTest, Int8ConstOpSizeInBytes) {
  constexpr absl::string_view kConstOpExpr =
      R"mlir(%cst = "tf.Const"() {value = dense<2> : tensor<3xi8>} : () -> tensor<3xi8>)mlir";

  Block block{};
  TF::ConstOp int_tensor_const_op = ParseConstOp(kConstOpExpr, block, ctx_);

  const int64_t num_bytes = GetSizeInBytes(int_tensor_const_op);
  EXPECT_THAT(num_bytes, Eq(3));
}

func.func @main(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tfl.less"(%arg0, %arg1) : (tensor<1xf32>, tensor<1xf32>) -> tensor<1xi1>
  %1 = "tf.If"(%0, %arg0, %arg1) {else_branch = @cond_false, then_branch = @cond_true, is_stateless = false} : (tensor<1xi1>, tensor<1xf32>, tensor<1xf32>) -> tensor<1xf32>
  func.return %1 : tensor<1xf32>
}

// CHECK-NEXT: %[[V3:.*]] = "tf.PartitionedCall"(%[[V2]], %[[CST_4]], %[[CST_5]]) <{config = "", config_proto = "", executor_type = "", f = @dequantize_i8}> : (tensor<*xi8>, tensor<f32>, tensor<i32>) -> tensor<*xf32>
// CHECK-NEXT: return %[[V3]] : tensor<*xf32>

// CHECK: func private @quantize_i8(

  // CHECK-NEXT:  buffer: 4
  // CHECK-NEXT:  } ]
  // CHECK-NEXT:  signature_defs: [ ]
  // CHECK-NEXT:}

  %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x!quant.uniform<i8:f32, 0.1>>, value = dense<2> : tensor<3xi8>} : () -> tensor<3x!quant.uniform<i8:f32, 0.1>>

  // CHECK-NEXT:  buffer: 4
  // CHECK-NEXT:  } ]
  // CHECK-NEXT:  signature_defs: [ ]
  // CHECK-NEXT:}

  %0 = "tfl.pseudo_qconst"() { qtype = tensor<3x!quant.uniform<i8:f32, 1.0>>, value = dense<2> : tensor<3xi8>} : () -> tensor<3x!quant.uniform<i8:f32, 1.0>>