EXPECT_THAT(num_bytes, Eq(16));
}

TEST_F(GetSizeInBytesTest, Bfloat16ConstOpSizeInBytes) {
  constexpr absl::string_view kConstOpExpr = R"mlir(
    %cst = "tf.Const"() {value = dense<1.0> : tensor<7xbf16>} : () -> tensor<7xbf16>
  )mlir";

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

namespace mlir {
namespace quant {

// Returns the size in bytes of the underlying data of `const_op`. If the
// underlying type's size cannot be determined, it assumes 4 bytes per element.
int64_t GetSizeInBytes(TF::ConstOp const_op);

}  // namespace quant
}  // namespace mlir

// `kAssumedNumBytesPerElem` bytes.
int64_t GetSizeOfUnsupportedTypeConst(TF::ConstOp const_op) {
  return kAssumedNumBytesPerElem * const_op.getValue().getNumElements();
}

}  // namespace

int64_t GetSizeInBytes(TF::ConstOp const_op) {
  const Type dtype = const_op.getDtype();

  if (dtype.isIntOrFloat()) {
    return GetSizeOfIntOrFloatConst(const_op);
  } else if (isa<TF::StringType>(dtype)) {

    absl::c_copy_if(func_op.getOps<TF::ConstOp>(),
                    std::back_inserter(target_const_ops),
                    [size_threshold](TF::ConstOp const_op) -> bool {
                      return GetSizeInBytes(const_op) > size_threshold;
                    });
  }

  return target_const_ops;
}

// Replaces every uses of ConstOps in `target_const_ops` to VarHandleOp ->