Open ilan-gold opened 4 days ago
@ilan-gold This is a good discussion worth having. #1967 was meant to be a first pass at adding device support and it exposed many areas of the code that were relying implicitly on numpy and CPU backed arrays in general. In https://github.com/zarr-developers/zarr-python/issues/1751, I proposed supporting both the __cuda_array_interface__
as well as dlpack but left out the dlpack support in #1967 mainly for simplicity. It would certainly be useful to add a dlpack based Buffer
implementation.
I was not aware of array_api
and it does seem like a really useful standard that we can use for a generic Buffer
implementation. The entrypoints based autodiscovery is a useful feature as well. From a quick search, here are the libraries that currently do/do not support the array_api
specification:
Given the current state of array library compliance, might it make sense to add an array_api
based Buffer
implementation but keep the existing explicit NumPy and CuPy based implementations with a view to move to the array_api
based implementation down the line?
@akshaysubr I would say that sounds fine, but perhaps not keep the explicit numpy
one since it is officially supported as array_api
. This would just leave cupy
for which there is this: https://github.com/data-apis/array-api-compat if I am not mistaken. I think the goal should really just be one implementation. I imagine it's doable modulo one or two extra if-then checks for things that might be broken. In fact, it seems like we won't have to wait too long if we don't want the compat layer: https://github.com/cupy/cupy/issues/8470
I really think we should strive for one implementation - simpler, and if there are bugs, maybe we can upstream the fixes since array api should be a one-stop-shop!
Also re: pytorch, I think it's ok if zarr
doesn't support it at first. This could be a good impetus to get its developers back on board re: array_api
. Also the data interchange with cupy is quite good and anyone using zarr
with torch for deep learning would probably have to write a data loader of some sort, no? So custom code already there as a layer around every read. Just some quick thoughts on that, but not as familiar.
Alternatively we could make the buffer class customizable to the point where you have an array-api compliant array as the "base" buffer, and then we decorate all functions that return arrays to return a dlpack
representation in a different, less-compliant library i.e., cupy as base and then torch on top. Same thing would probably enable tensorflow (experimentally, maybe, just from a quick google): https://www.tensorflow.org/api_docs/python/tf/experimental/dlpack/from_dlpack + https://docs.cupy.dev/en/stable/user_guide/interoperability.html#dlpack
Discussed in https://github.com/zarr-developers/zarr-python/discussions/2197
cc: @akshaysubr @madsbk @d-v-b @dcherian