bfloat16 unsupported by torch.tensor that takes np arrays inputs

[Ishticode]

I am not sure whether to raise this here or at pytorch but here it is.
Usually this works works

x = np.array([2.], dtype=np.float16)
x = torch.tensor(x)

but not

from ml_dtypes import bfloat16
x = np.array([2.], dtype=bfloat16)
x = torch.tensor(x)

resulting in error

TypeError: can't convert np.ndarray of type bfloat16. The only supported types are: float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.

although pytorch does support bfloat16.
What do you say? or should I take this to PyTorch. Thank you.

Note: tensorflow does accept np arrays of ml_dtypes extended bfloat16 without any issue.
Note: I think u will be able to reproduce with almost all versions. But I can give the specifics if versions have anything to do with it.

[jakevdp]

Thanks for the report! This is something that would have to be addressed by the torch project. I'd suggest raising the issue there.

Recent versions of jax and tensorflow explicitly depend on the ml_dtypes project, and are built to recognize the types defined here. PyTorch could certainly do the same if it aligns with the goals of the developers of that project, but it can't happen without that code being added to the pytorch source.

[hawkinsp]

Closing, because this is a pytorch issue.

[ariG23498]

Hey folks!
Sorry to ping on a closed issue 😅

I wonder if there was a way to port torch tensors (bfloat16) to jax Array?
My workflow currently is:

torch_tensor = torch.randn((1, 2), dtype=torch.bfloat16)
intermediate_tensor = torch_tensor.float()    # here torch_tensor.numpy() fails
jax_tensor = jnp.array(intermediate_tensor)

I would like to do away with the increased precision (float32) conversion.
My apologies if this is already solved or I am missing something (a link to the correct resource would be highly appreciated!)

[jakevdp]

I would probably do something like this:

torch_tensor = torch.randn((1, 2), dtype=torch.bfloat16)
intermediate_tensor = torch_tensor.view(torch.uint16)
jax_tensor = jnp.array(intermediate_tensor).view('bfloat16')

It's really a shame that torch uses a bfloat16 implementation that's incompatible with NumPy.

[hawkinsp]

You can also exchange data with DLPack, which should support BF16 irrespective of any NumPy support.