python - tensorflow:注册 numpy bfloat16 扩展

Question

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

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'

有没有人看到我做错了什么？或者如何启用这个 numpy 扩展？

非常感谢

Answer 1

试试这个:

from tensorflow.python import _pywrap_bfloat16


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

print(bfloat16(1.0))
# bfloat16(1)