%arg2: tensor<5xf32> {tf._user_specified_name = "b1"},
                             %arg3: tensor<3920x10xf32> {tf._user_specified_name = "w2"},
                             %arg4: tensor<10xf32> {tf._user_specified_name = "b2"}) -> tensor<10xf32>

  EXPECT_THAT(num_bytes, Eq(3));
}

TEST_F(GetSizeInBytesTest, Float32ConstOpSizeInBytes) {
  constexpr absl::string_view kConstOpExpr =
      R"mlir(%cst = "tf.Const"() {value = dense<3.0> : tensor<4xf32>} : () -> tensor<4xf32>)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(16));
}

    } : (tensor<*xf32>, tensor<144x10xf32>, tensor<10xf32>) -> tensor<*xf32>
    %5 = "quantfork.qcast"(%4) : (tensor<*xf32>) -> tensor<*x!quant.uniform<i8:f32, 0.1>>
    %6 = "quantfork.dcast"(%5) : (tensor<*x!quant.uniform<i8:f32, 0.1>>) -> tensor<*xf32>
    func.return %6 : tensor<*xf32>
  }

  func.func private @composite_matmul_with_bias_fn_1(%a: tensor<*xf32>, %b: tensor<*xf32>, %c: tensor<*xf32>) -> tensor<*xf32> {

func.func @tfl_wrapped_jax_random_normal(%arg0: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
  // This is a fake jax random normal body.
  %0 = stablehlo.constant dense<0.0> : tensor<12xf32>
  %1 = "stablehlo.reshape"(%0) : (tensor<12xf32>) -> tensor<3x4xf32>
  %2 = "stablehlo.tuple"(%1) : (tensor<3x4xf32>) -> tuple<tensor<3x4xf32>>
  func.return %2 : tuple<tensor<3x4xf32>>
}


// CHECK-LABEL:   func @tfl_wrapped_jax_random_uniform(

  func.func @matmul_with_reshape_disable_quantization(%arg0: tensor<1x10xf32>, %arg1: tensor<10x10xf32>) -> (tensor<?x10xf32>) {
    %cst = "tf.Const"() {value = dense<0.000000e+00> : tensor<10xf32>} : () -> tensor<10xf32>
    %cst_0 = "tf.Const"() {value = dense<[-1, 10]> : tensor<2xi32>} : () -> tensor<2xi32>
    %1 = "tf.MatMul"(%arg0, %arg1) {
      transpose_a = false, transpose_b = false

  %0 = mhlo.constant dense<0.000000e+00> : tensor<32x10xf32>
  %1 = mhlo.constant dense<0.000000e+00> : tensor<32x28x28x5xf32>
  %2 = "tf.ReadVariableOp"(%arg4) : (tensor<!tf_type.resource<tensor<10xf32>>>) -> tensor<10xf32>

  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;

    %4 = "quantfork.stats"(%2) {layerStats = dense<[-2795.36523, 4609.57373]> : tensor<2xf32>} : (tensor<*xf32>) -> tensor<*xf32>
    %5 = "tf.BiasAdd"(%4, %3) {data_format = "NHWC", device = ""} : (tensor<*xf32>, tensor<2xf32>) -> tensor<*xf32>
    %6 = "quantfork.stats"(%5) {layerStats = dense<[-2788.31055, 4616.62842]> : tensor<2xf32>} : (tensor<*xf32>) -> tensor<*xf32>