Kevin Meier
Kevin Meier

Reputation: 2582

tensorflow: register numpy bfloat16 extension

As I've seen there is a numpy extension for bfloat16 in tensorflow:

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

This extension can be enabled by calling RegisterNumpyBfloat16. Or at least it should. I've installed tensorflow 2.4.1 and tried to enable the extension, but I just get this error:

> tf.RegisterNumpyBfloat16()

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-6-4f89beb80796> in <module>
----> 1 tf.RegisterNumpyBfloat16()

AttributeError: module 'tensorflow' has no attribute 'RegisterNumpyBfloat16'

Does anyone see what I do wrong? Or how can this numpy-extension be enabled?

Thank you very much

Upvotes: 3

Views: 4880

Answers (2)

James Mishra
James Mishra

Reputation: 4510

I just tried this on TensorFlow 2.4.0 and 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)

Upvotes: 5

o-90
o-90

Reputation: 17603

Try this:

from tensorflow.python import _pywrap_bfloat16


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

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

Upvotes: 3

Related Questions