Boolean DeviceArray to DLPack

[yuanqing-wang]

Hi it seems that Boolean DeviceArray cannot be converted to DLPack?

>>> import jax
>>> from jax import numpy as jnp
>>> x = jnp.array([True, False])
>>> from jax import dlpack
>>> dlpack.to_dlpack(x)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/Users/wangy1/Documents/GitHub/jax/jax/_src/dlpack.py", line 42, in to_dlpack
    buf, take_ownership=take_ownership)
RuntimeError: Unimplemented: XLA type PRED has no DLPack equivalent

[hawkinsp]

That's because DLPack has no well-defined way to represent "bool"s:

https://github.com/dmlc/dlpack/blob/1b794e7088b754f4b0398d211452de0ab28312b5/include/dlpack/dlpack.h#L77

I could choose an encoding arbitrarily (e.g., as int8s), but that may or may not match any other system you want to share data with, and you could also achieve that by first converting to int8 and exporting that array instead.

I'm happy to add a DLPack bool export if we can agree on form it should take...

[yuanqing-wang]

Looks like pytorch does convert Booleans to int8s.

>>> import torch
>>> from torch.utils import dlpack
>>> dlpack.from_dlpack(dlpack.to_dlpack(torch.tensor([True, False])))
tensor([1, 0], dtype=torch.uint8)

[hawkinsp]

My temptation would be not to do this cast automatically: in some sense, it seems better to me to be explicit and have the user cast to uint8 if that's what they want. Would it help if we improved the error message a bit to suggest this explicitly?

[hawkinsp]

This is now fixed at head. We allow the export of bool arrays as uint8, same as PyTorch. You can't import bool arrays (since there is no such thing as a DLPack bool array at the moment).