add std::ostream& operator<< for BFloat16 in BFloat16.h (pytorch#121302)

This PR Move `operator<<` of `BFloat16` to `BFloat16.h`.

Previously, this function is in `TensorDataContainer.h`. If need `std::cout` a `BFloat16` variable when debugging, `TensorDataContainer.h` have to be included. This is inconvient and counterintuitive.

Other dtypes such as `Half`, define their `operator<<` in headers where they are defined such as `Half.h`. Therefore, I think it makes more sense to move `operator<<` of `BFloat16` to `BFloat16.h`

Pull Request resolved: pytorch#121302
Approved by: https://github.com/ezyang

Author: LamForest

c10/util/BFloat16.h

@@ -8,6 +8,8 @@
 #include <cstdint>
 #include <cstring>
 
+#include <iosfwd>
+
 #if defined(__CUDACC__) && !defined(USE_ROCM)
 #include <cuda_bf16.h>
 #endif
@@ -112,6 +114,8 @@ struct alignas(2) BFloat16 {
 #endif
 };
 
+C10_API std::ostream& operator<<(std::ostream& out, const BFloat16& value);
+
 } // namespace c10
 
 #include <c10/util/BFloat16-inl.h> // IWYU pragma: keep

c10/util/Bfloat16.cpp

@@ -0,0 +1,15 @@
+#include <c10/util/BFloat16.h>
+#include <ostream>
+#include <type_traits>
+
+namespace c10 {
+
+static_assert(
+    std::is_standard_layout_v<BFloat16>,
+    "c10::BFloat16 must be standard layout.");
+
+std::ostream& operator<<(std::ostream& out, const BFloat16& value) {
+  out << (float)value;
+  return out;
+}
+} // namespace c10

c10/util/Float8_e4m3fn.cpp

@@ -1,5 +1,5 @@
 #include <c10/util/Float8_e4m3fn.h>
-#include <iostream>
+#include <ostream>
 #include <type_traits>
 
 namespace c10 {

c10/util/Float8_e4m3fnuz.cpp

@@ -1,5 +1,5 @@
 #include <c10/util/Float8_e4m3fnuz.h>
-#include <iostream>
+#include <ostream>
 
 namespace c10 {
 

c10/util/Float8_e5m2.cpp

@@ -1,5 +1,5 @@
 #include <c10/util/Float8_e5m2.h>
-#include <iostream>
+#include <ostream>
 
 namespace c10 {
 

c10/util/Float8_e5m2fnuz.cpp

@@ -1,5 +1,5 @@
 #include <c10/util/Float8_e5m2fnuz.h>
-#include <iostream>
+#include <ostream>
 
 namespace c10 {
 

c10/util/Half.cpp

@@ -1,5 +1,5 @@
 #include <c10/util/Half.h>
-#include <iostream>
+#include <ostream>
 #include <type_traits>
 
 namespace c10 {

torch/csrc/api/include/torch/detail/TensorDataContainer.h

@@ -28,15 +28,6 @@ inline std::ostream& operator<<(
     std::ostream& stream,
     const TensorDataContainer& tensor_data_container);
 
-// FIXME: There is no `operator<<` overload for `at::kBFloat16` type,
-// and we need to convert it to `float` type using `operator float()` function
-// defined in `c10/util/BFloat16.h`.
-// Tracking issue: https://github.com/pytorch/pytorch/issues/28845
-inline std::ostream& operator<<(std::ostream& stream, c10::BFloat16 value) {
-  stream << static_cast<float>(value);
-  return stream;
-}
-
 inline c10::ScalarType compute_desired_dtype(c10::ScalarType scalar_type) {
   if (scalar_type == at::kInt || scalar_type == at::kLong) {
     // C++ `torch::tensor` with an integer type or an `at::ArrayRef` /