pytorch

fa623086 - [tensorboard] Fix TensorBoard summary encoding for torch.bfloat16 tensors (#108351)

[tensorboard] Fix TensorBoard summary encoding for torch.bfloat16 tensors (#108351) Summary: The `tensor_proto` function in the TensorBoard summary writer code doesn't correctly encode `torch.bfloat16` tensors; it tries to use a data type of `DT_BFLOAT` when creating the protobuf, but `DT_BFLOAT` is not a valid enum value (see `types.proto`). The correct value to use when encoding tensors of this type is `DT_BLOAT16`. This diff updates the type map in the summary code to use the correct type. While fixing this error, I also noticed the wrong field of the protobuf was being used when encoding tensors of this type; per the docs in the proto file, the DT_HALF and DT_BFLOAT16 types should use the `half_val` field, not `float_val`. Since this might confuse folks trying to read this data from storage in the future, I've updated the code to correctly use to `half_val` field for these cases. Note that there's no real size advantage from doing this, since both the `half_val` and `float_val` fields are 32 bits long. Test Plan: Added a parameterized unit test that tests encoding tensors with `torch.half`, `torch.float16`, and `torch.bfloat16` data types. # Before this change The test fails with an `ValueError` due to the incorrect enum label: ``` ====================================================================== ERROR: test_bfloat16_tensor_proto (test_tensorboard.TestTensorProtoSummary) ---------------------------------------------------------------------- Traceback (most recent call last): File "/data/users/jcarreiro/fbsource/buck-out/v2/gen/fbcode/f88b3f368c9334db/caffe2/test/__tensorboard__/tensorboard#link-tree/torch/testing/_internal/common_utils.py", line 2382, in wrapper method(*args, **kwargs) File "/data/users/jcarreiro/fbsource/buck-out/v2/gen/fbcode/f88b3f368c9334db/caffe2/test/__tensorboard__/tensorboard#link-tree/test_tensorboard.py", line 871, in test_bfloat16_tensor_proto tensor_proto( File "/data/users/jcarreiro/fbsource/buck-out/v2/gen/fbcode/f88b3f368c9334db/caffe2/test/__tensorboard__/tensorboard#link-tree/torch/utils/tensorboard/summary.py", line 400, in tensor_proto tensor_proto = TensorProto(**tensor_proto_args) ValueError: unknown enum label "DT_BFLOAT" To execute this test, run the following from the base repo dir: python test/__tensorboard__/tensorboard#link-tree/test_tensorboard.py -k test_bfloat16_tensor_proto This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0 ---------------------------------------------------------------------- ``` # After this change The test passes. Reviewed By: tanvigupta17 Differential Revision: D48828958 Pull Request resolved: https://github.com/pytorch/pytorch/pull/108351 Approved by: https://github.com/hamzajzmati, https://github.com/XilunWu