How to enable bfloat16 PyTorch/XLA on Google Colab

sakaia · March 12, 2019, 12:05am

In PyTorch/XLA implementation, an environment vatiable XLA_USE_BF16 exists.
Regrettably It cannot see the TF_LOG on Google Colab.
(If some one know, how to set TF_LOG enable, please notice me.)
Is there any way to enable bfloat16 for Pytorch/XLA on Google Colab?

Ref.

github.com

pytorch/xla/blob/master/torch_xla/csrc/tensor_util.cpp#L21


#include "tensorflow/compiler/xla/xla_client/sys_util.h"
#include "tensorflow/compiler/xla/xla_client/tf_logging.h"
#include "tensorflow/compiler/xla/xla_client/util.h"
#include "tensorflow/core/lib/bfloat16/bfloat16.h"
#include "torch_xla/csrc/helpers.h"


namespace torch_xla {
namespace {


bool ShouldUseBF16() {
  int use_fp16 = xla::sys_util::GetEnvInt("XLA_USE_BF16", 0);
  if (use_fp16 != 0) {
    TF_LOG(INFO) << "Using BF16 data type for floating point values";
  }
  return use_fp16 != 0;
}


xla::PrimitiveType XlaTypeFromTensorType(at::ScalarType scalar_type,
                                         const Device& device) {
  switch (scalar_type) {
    case at::ScalarType::Double:

sieginglion · March 14, 2020, 6:51am

try %env XLA_USE_BF16=1