ENH: Support DLPack for scalars by converting to 0-D array#31893
ENH: Support DLPack for scalars by converting to 0-D array#31893MaanasArora wants to merge 3 commits into
Conversation
a7f6f8d to
0c7fab4
Compare
0c7fab4 to
de0d3de
Compare
|
I Just pushed an update to reject |
|
Added myself to the reviewers here as its one of the things I wanted to review when I find the next time slot. |
ikrommyd
left a comment
There was a problem hiding this comment.
I think this is good and clean! Only left two minor comments.
| if (result_device.device_type != kDLCPU || | ||
| result_device.device_id != 0) { | ||
| PyErr_SetString(PyExc_BufferError, | ||
| "Cannot export a scalar to DLPack on a non-CPU device."); | ||
| return NULL; | ||
| } |
There was a problem hiding this comment.
I think this is dead code. device_converter should already guarantee that device_type will be kDLCPU. Unless you wanted it for a different error message.
| if (result_device.device_type != kDLCPU || | |
| result_device.device_id != 0) { | |
| PyErr_SetString(PyExc_BufferError, | |
| "Cannot export a scalar to DLPack on a non-CPU device."); | |
| return NULL; | |
| } |
There was a problem hiding this comment.
I'd add testing for the dl_device argument and for the copy=True case which if I'm seeing correctly are currently missing.
|
@seberg based on your comment #27137 (comment) In [1]: import ml_dtypes
In [2]: ml_dtypes.bfloat16(4).__dlpack__()
---------------------------------------------------------------------------
BufferError Traceback (most recent call last)
Cell In[2], line 1
----> 1 ml_dtypes.bfloat16(4).__dlpack__()
BufferError: DLPack only supports signed/unsigned integers, float and complex dtypes (or dtypes registered by third-party packages).
In [3]: from numpy_quaddtype import QuadPrecDType, QuadPrecision
In [4]: QuadPrecision(4).__dlpack__()
---------------------------------------------------------------------------
BufferError Traceback (most recent call last)
Cell In[4], line 1
----> 1 QuadPrecision(4).__dlpack__()
BufferError: DLPack only supports signed/unsigned integers, float and complex dtypes (or dtypes registered by third-party packages). |
Closes #27137 by adding
__dlpack__and__dlpack_device__methods to NumPy scalars.scalar.__dlpack__currently delegates tondarray.__dlpack__by converting the scalar to a 0-D array. This is a very minimal solution - a copy is fairly natural in this context, so this workaround seemed natural as well. ping @seberg - thanks!I have to check if we can add additional handling for the
copy=kwarg - that seemed to be the main case where the array and scalar implementations might diverge in practice (we are always copying). Perhaps with certain conditions we can add validation forcopy. I also have to check if I missed some other nuances within the arguments, and needs a release note, but otherwise, this is ready for review!AI Disclosure
I used an agent to search for DLPack's implementation in more detail, to check for errors in my code, and to write a few tests (the
independent_copyone, mainly).