tensorflow:注册numpy bfloat16扩展

Question

正如我所看到的，在tensorflow中有一个numpy扩展bfloat16:

https://github.com/tensorflow/tensorflow/blob/24ffe9f729160a095a5cab8f5923920182803182/tensorflow/python/lib/core/bfloat16_wrapper.cc

这个扩展可以通过调用RegisterNumpyBfloat16来启用。或者至少应该如此。我已经安装了tensorflow 2.4.1，并试图启用扩展，但我只是得到这个错误:

> tf.RegisterNumpyBfloat16()
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-6-4f89beb80796> in <module>
----> 1 tf.RegisterNumpyBfloat16()
AttributeError: module 'tensorflow' has no attribute 'RegisterNumpyBfloat16'

有人看到我做错了什么吗?或者如何启用这个愚蠢的扩展?

Thank you very much

Answer 1

我刚刚在TensorFlow 2.4.0和NumPy 1.19.4上尝试了这个。

import numpy as np
import tensorflow as tf
bfloat16 = tf.bfloat16.as_numpy_dtype
np.array([1.0, 2.0, 3.0], dtype=bfloat16)
# array([bfloat16(1), bfloat16(2), bfloat16(3)], dtype=bfloat16)

Answer 2

试试这个:

from tensorflow.python import _pywrap_bfloat16

bfloat16 = _pywrap_bfloat16.TF_bfloat16_type()
print(bfloat16)
# <class 'bfloat16'>
print(bfloat16(1.0))
# bfloat16(1)